add xpu and npu support for ernie series. (#1639)

modules/text/text_generation/ernie_gen_acrostic_poetry/module.py

@@ -60,8 +60,36 @@ class ErnieGen(hub.NLPPredictionModule):
         self.rev_dict[self.tokenizer.vocab['[UNK]']] = ''  # replace [PAD]
         self.rev_lookup = np.vectorize(lambda i: self.rev_dict[i])
 
+        # detect npu
+        npu_id = self._get_device_id("FLAGS_selected_npus")
+        if npu_id != -1:
+            # use npu
+            self.use_device = "npu"
+        else:
+            # detect gpu
+            gpu_id = self._get_device_id("CUDA_VISIBLE_DEVICES")
+            if gpu_id != -1:
+                # use gpu
+                self.use_device = "gpu"
+            else:
+                # detect xpu
+                xpu_id = self._get_device_id("XPU_VISIBLE_DEVICES")
+                if xpu_id != -1:
+                    # use xpu
+                    self.use_device = "xpu"
+                else:
+                    self.use_device = "cpu"
+
+    def _get_device_id(self, places):
+        try:
+            places = os.environ[places]
+            id = int(places)
+        except:
+            id = -1
+        return id
+
     @serving
-    def generate(self, texts, use_gpu=False, beam_width=5):
+    def generate(self, texts, use_gpu=False, beam_width=5, use_device=None):
         """
         Get the continuation of the input poetry.
 
@@ -69,6 +97,7 @@ class ErnieGen(hub.NLPPredictionModule):
              texts(list): the front part of a poetry.
              use_gpu(bool): whether use gpu to predict or not
              beam_width(int): the beam search width.
+             use_device (str): use cpu, gpu, xpu or npu, overwrites use_gpu flag.
 
         Returns:
              results(list): the poetry continuations.
@@ -91,13 +120,25 @@ class ErnieGen(hub.NLPPredictionModule):
                         'The input text: %s, contains non-Chinese characters, which may result in magic output' % text)
                     break
 
-        if use_gpu and "CUDA_VISIBLE_DEVICES" not in os.environ:
-            use_gpu = False
-            logger.warning(
-                "use_gpu has been set False as you didn't set the environment variable CUDA_VISIBLE_DEVICES while using use_gpu=True"
-            )
+        if use_device is not None:
+            # check 'use_device' match 'device on init'
+            if use_device != self.use_device:
+                raise RuntimeError(
+                    "the 'use_device' parameter when calling generate, does not match internal device found on init.")
+        else:
+            # use_device is None, follow use_gpu flag
+            if use_gpu == False:
+                use_device = "cpu"
+            elif use_gpu == True and self.use_device != 'gpu':
+                use_device = "cpu"
+                logger.warning(
+                    "use_gpu has been set False as you didn't set the environment variable CUDA_VISIBLE_DEVICES while using use_gpu=True"
+                )
+            else:
+                # use_gpu and self.use_device are both true
+                use_device = "gpu"
 
-        paddle.set_device('gpu') if use_gpu else paddle.set_device('cpu')
+        paddle.set_device(use_device)
 
         self.model.eval()
         results = []
@@ -135,8 +176,11 @@ class ErnieGen(hub.NLPPredictionModule):
         """
         self.arg_config_group.add_argument(
             '--use_gpu', type=ast.literal_eval, default=False, help="whether use GPU for prediction")
-
         self.arg_config_group.add_argument('--beam_width', type=int, default=5, help="the beam search width")
+        self.arg_config_group.add_argument(
+            '--use_device',
+            choices=["cpu", "gpu", "xpu", "npu"],
+            help="use cpu, gpu, xpu or npu. overwrites use_gpu flag.")
 
     @runnable
     def run_cmd(self, argvs):
@@ -164,7 +208,8 @@ class ErnieGen(hub.NLPPredictionModule):
             self.parser.print_help()
             return None
 
-        results = self.generate(texts=input_data, use_gpu=args.use_gpu, beam_width=args.beam_width)
+        results = self.generate(
+            texts=input_data, use_gpu=args.use_gpu, beam_width=args.beam_width, use_device=args.use_device)
 
         return results
 

modules/text/text_generation/ernie_gen_couplet/module.py

@@ -54,8 +54,36 @@ class ErnieGen(hub.NLPPredictionModule):
         self.rev_dict[self.tokenizer.vocab['[UNK]']] = ''  # replace [PAD]
         self.rev_lookup = np.vectorize(lambda i: self.rev_dict[i])
 
+        # detect npu
+        npu_id = self._get_device_id("FLAGS_selected_npus")
+        if npu_id != -1:
+            # use npu
+            self.use_device = "npu"
+        else:
+            # detect gpu
+            gpu_id = self._get_device_id("CUDA_VISIBLE_DEVICES")
+            if gpu_id != -1:
+                # use gpu
+                self.use_device = "gpu"
+            else:
+                # detect xpu
+                xpu_id = self._get_device_id("XPU_VISIBLE_DEVICES")
+                if xpu_id != -1:
+                    # use xpu
+                    self.use_device = "xpu"
+                else:
+                    self.use_device = "cpu"
+
+    def _get_device_id(self, places):
+        try:
+            places = os.environ[places]
+            id = int(places)
+        except:
+            id = -1
+        return id
+
     @serving
-    def generate(self, texts, use_gpu=False, beam_width=5):
+    def generate(self, texts, use_gpu=False, beam_width=5, use_device=None):
         """
         Get the right rolls from the left rolls.
 
@@ -63,6 +91,7 @@ class ErnieGen(hub.NLPPredictionModule):
              texts(list): the left rolls.
              use_gpu(bool): whether use gpu to predict or not
              beam_width(int): the beam search width.
+             use_device (str): use cpu, gpu, xpu or npu, overwrites use_gpu flag.
 
         Returns:
              results(list): the right rolls.
@@ -80,13 +109,25 @@ class ErnieGen(hub.NLPPredictionModule):
                         'The input text: %s, contains non-Chinese characters, which may result in magic output' % text)
                     break
 
-        if use_gpu and "CUDA_VISIBLE_DEVICES" not in os.environ:
-            use_gpu = False
-            logger.warning(
-                "use_gpu has been set False as you didn't set the environment variable CUDA_VISIBLE_DEVICES while using use_gpu=True"
-            )
-
-        paddle.set_device('gpu') if use_gpu else paddle.set_device('cpu')
+        if use_device is not None:
+            # check 'use_device' match 'device on init'
+            if use_device != self.use_device:
+                raise RuntimeError(
+                    "the 'use_device' parameter when calling generate, does not match internal device found on init.")
+        else:
+            # use_device is None, follow use_gpu flag
+            if use_gpu == False:
+                use_device = "cpu"
+            elif use_gpu == True and self.use_device != 'gpu':
+                use_device = "cpu"
+                logger.warning(
+                    "use_gpu has been set False as you didn't set the environment variable CUDA_VISIBLE_DEVICES while using use_gpu=True"
+                )
+            else:
+                # use_gpu and self.use_device are both true
+                use_device = "gpu"
+
+        paddle.set_device(use_device)
 
         self.model.eval()
         results = []
@@ -124,8 +165,11 @@ class ErnieGen(hub.NLPPredictionModule):
         """
         self.arg_config_group.add_argument(
             '--use_gpu', type=ast.literal_eval, default=False, help="whether use GPU for prediction")
-
         self.arg_config_group.add_argument('--beam_width', type=int, default=5, help="the beam search width")
+        self.arg_config_group.add_argument(
+            '--use_device',
+            choices=["cpu", "gpu", "xpu", "npu"],
+            help="use cpu, gpu, xpu or npu. overwrites use_gpu flag.")
 
     @runnable
     def run_cmd(self, argvs):
@@ -153,7 +197,8 @@ class ErnieGen(hub.NLPPredictionModule):
             self.parser.print_help()
             return None
 
-        results = self.generate(texts=input_data, use_gpu=args.use_gpu, beam_width=args.beam_width)
+        results = self.generate(
+            texts=input_data, use_gpu=args.use_gpu, beam_width=args.beam_width, use_device=args.use_device)
 
         return results
 

modules/text/text_generation/ernie_gen_lover_words/module.py

@@ -54,8 +54,36 @@ class ErnieGen(hub.NLPPredictionModule):
         self.rev_dict[self.tokenizer.vocab['[UNK]']] = ''  # replace [PAD]
         self.rev_lookup = np.vectorize(lambda i: self.rev_dict[i])
 
+        # detect npu
+        npu_id = self._get_device_id("FLAGS_selected_npus")
+        if npu_id != -1:
+            # use npu
+            self.use_device = "npu"
+        else:
+            # detect gpu
+            gpu_id = self._get_device_id("CUDA_VISIBLE_DEVICES")
+            if gpu_id != -1:
+                # use gpu
+                self.use_device = "gpu"
+            else:
+                # detect xpu
+                xpu_id = self._get_device_id("XPU_VISIBLE_DEVICES")
+                if xpu_id != -1:
+                    # use xpu
+                    self.use_device = "xpu"
+                else:
+                    self.use_device = "cpu"
+
+    def _get_device_id(self, places):
+        try:
+            places = os.environ[places]
+            id = int(places)
+        except:
+            id = -1
+        return id
+
     @serving
-    def generate(self, texts, use_gpu=False, beam_width=5):
+    def generate(self, texts, use_gpu=False, beam_width=5, use_device=None):
         """
         Get the continuation of the input poetry.
 
@@ -63,6 +91,7 @@ class ErnieGen(hub.NLPPredictionModule):
              texts(list): the front part of a poetry.
              use_gpu(bool): whether use gpu to predict or not
              beam_width(int): the beam search width.
+             use_device (str): use cpu, gpu, xpu or npu, overwrites use_gpu flag.
 
         Returns:
              results(list): the poetry continuations.
@@ -74,12 +103,25 @@ class ErnieGen(hub.NLPPredictionModule):
         else:
             raise ValueError("The input texts should be a list with nonempty string elements.")
 
-        if use_gpu and "CUDA_VISIBLE_DEVICES" not in os.environ:
-            use_gpu = False
-            logger.warning(
-                "use_gpu has been set False as you didn't set the environment variable CUDA_VISIBLE_DEVICES while using use_gpu=True"
-            )
-        paddle.set_device('gpu') if use_gpu else paddle.set_device('cpu')
+        if use_device is not None:
+            # check 'use_device' match 'device on init'
+            if use_device != self.use_device:
+                raise RuntimeError(
+                    "the 'use_device' parameter when calling generate, does not match internal device found on init.")
+        else:
+            # use_device is None, follow use_gpu flag
+            if use_gpu == False:
+                use_device = "cpu"
+            elif use_gpu == True and self.use_device != 'gpu':
+                use_device = "cpu"
+                logger.warning(
+                    "use_gpu has been set False as you didn't set the environment variable CUDA_VISIBLE_DEVICES while using use_gpu=True"
+                )
+            else:
+                # use_gpu and self.use_device are both true
+                use_device = "gpu"
+
+        paddle.set_device(use_device)
         self.model.eval()
         results = []
         for text in predicted_data:
@@ -116,8 +158,11 @@ class ErnieGen(hub.NLPPredictionModule):
         """
         self.arg_config_group.add_argument(
             '--use_gpu', type=ast.literal_eval, default=False, help="whether use GPU for prediction")
-
         self.arg_config_group.add_argument('--beam_width', type=int, default=5, help="the beam search width")
+        self.arg_config_group.add_argument(
+            '--use_device',
+            choices=["cpu", "gpu", "xpu", "npu"],
+            help="use cpu, gpu, xpu or npu. overwrites use_gpu flag.")
 
     @runnable
     def run_cmd(self, argvs):
@@ -145,7 +190,8 @@ class ErnieGen(hub.NLPPredictionModule):
             self.parser.print_help()
             return None
 
-        results = self.generate(texts=input_data, use_gpu=args.use_gpu, beam_width=args.beam_width)
+        results = self.generate(
+            texts=input_data, use_gpu=args.use_gpu, beam_width=args.beam_width, use_device=args.use_device)
 
         return results
 

modules/text/text_generation/ernie_gen_poetry/module.py

@@ -54,8 +54,36 @@ class ErnieGen(hub.NLPPredictionModule):
         self.rev_dict[self.tokenizer.vocab['[UNK]']] = ''  # replace [PAD]
         self.rev_lookup = np.vectorize(lambda i: self.rev_dict[i])
 
+        # detect npu
+        npu_id = self._get_device_id("FLAGS_selected_npus")
+        if npu_id != -1:
+            # use npu
+            self.use_device = "npu"
+        else:
+            # detect gpu
+            gpu_id = self._get_device_id("CUDA_VISIBLE_DEVICES")
+            if gpu_id != -1:
+                # use gpu
+                self.use_device = "gpu"
+            else:
+                # detect xpu
+                xpu_id = self._get_device_id("XPU_VISIBLE_DEVICES")
+                if xpu_id != -1:
+                    # use xpu
+                    self.use_device = "xpu"
+                else:
+                    self.use_device = "cpu"
+
+    def _get_device_id(self, places):
+        try:
+            places = os.environ[places]
+            id = int(places)
+        except:
+            id = -1
+        return id
+
     @serving
-    def generate(self, texts, use_gpu=False, beam_width=5):
+    def generate(self, texts, use_gpu=False, beam_width=5, use_device=None):
         """
         Get the continuation of the input poetry.
 
@@ -63,6 +91,7 @@ class ErnieGen(hub.NLPPredictionModule):
              texts(list): the front part of a poetry.
              use_gpu(bool): whether use gpu to predict or not
              beam_width(int): the beam search width.
+             use_device (str): use cpu, gpu, xpu or npu, overwrites use_gpu flag.
 
         Returns:
              results(list): the poetry continuations.
@@ -91,12 +120,26 @@ class ErnieGen(hub.NLPPredictionModule):
                         % text)
                 break
 
-        if use_gpu and "CUDA_VISIBLE_DEVICES" not in os.environ:
-            use_gpu = False
-            logger.warning(
-                "use_gpu has been set False as you didn't set the environment variable CUDA_VISIBLE_DEVICES while using use_gpu=True"
-            )
-        paddle.set_device('gpu') if use_gpu else paddle.set_device('cpu')
+        if use_device is not None:
+            # check 'use_device' match 'device on init'
+            if use_device != self.use_device:
+                raise RuntimeError(
+                    "the 'use_device' parameter when calling generate, does not match internal device found on init.")
+        else:
+            # use_device is None, follow use_gpu flag
+            if use_gpu == False:
+                use_device = "cpu"
+            elif use_gpu == True and self.use_device != 'gpu':
+                use_device = "cpu"
+                logger.warning(
+                    "use_gpu has been set False as you didn't set the environment variable CUDA_VISIBLE_DEVICES while using use_gpu=True"
+                )
+            else:
+                # use_gpu and self.use_device are both true
+                use_device = "gpu"
+
+        paddle.set_device(use_device)
+
         self.model.eval()
         results = []
         for text in predicted_data:
@@ -133,8 +176,11 @@ class ErnieGen(hub.NLPPredictionModule):
         """
         self.arg_config_group.add_argument(
             '--use_gpu', type=ast.literal_eval, default=False, help="whether use GPU for prediction")
-
         self.arg_config_group.add_argument('--beam_width', type=int, default=5, help="the beam search width")
+        self.arg_config_group.add_argument(
+            '--use_device',
+            choices=["cpu", "gpu", "xpu", "npu"],
+            help="use cpu, gpu, xpu or npu. overwrites use_gpu flag.")
 
     @runnable
     def run_cmd(self, argvs):
@@ -162,7 +208,8 @@ class ErnieGen(hub.NLPPredictionModule):
             self.parser.print_help()
             return None
 
-        results = self.generate(texts=input_data, use_gpu=args.use_gpu, beam_width=args.beam_width)
+        results = self.generate(
+            texts=input_data, use_gpu=args.use_gpu, beam_width=args.beam_width, use_device=args.use_device)
 
         return results