From d79f6c49c6d8b58e37947ac0b4909b830897f9c8 Mon Sep 17 00:00:00 2001 From: houj04 <35131887+houj04@users.noreply.github.com> Date: Tue, 19 Oct 2021 16:25:05 +0800 Subject: [PATCH] add xpu and npu support for ernie series. (#1639) --- .../ernie_gen_acrostic_poetry/module.py | 63 +++++++++++++++--- .../ernie_gen_couplet/module.py | 65 ++++++++++++++++--- .../ernie_gen_lover_words/module.py | 64 +++++++++++++++--- .../ernie_gen_poetry/module.py | 65 ++++++++++++++++--- 4 files changed, 220 insertions(+), 37 deletions(-) diff --git a/modules/text/text_generation/ernie_gen_acrostic_poetry/module.py b/modules/text/text_generation/ernie_gen_acrostic_poetry/module.py index f39c0605..34d1e92b 100644 --- a/modules/text/text_generation/ernie_gen_acrostic_poetry/module.py +++ b/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 diff --git a/modules/text/text_generation/ernie_gen_couplet/module.py b/modules/text/text_generation/ernie_gen_couplet/module.py index 59bb39ad..4aaaba6e 100644 --- a/modules/text/text_generation/ernie_gen_couplet/module.py +++ b/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 diff --git a/modules/text/text_generation/ernie_gen_lover_words/module.py b/modules/text/text_generation/ernie_gen_lover_words/module.py index 55b12719..916857d3 100644 --- a/modules/text/text_generation/ernie_gen_lover_words/module.py +++ b/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 diff --git a/modules/text/text_generation/ernie_gen_poetry/module.py b/modules/text/text_generation/ernie_gen_poetry/module.py index a5b8f02c..ba7cea40 100644 --- a/modules/text/text_generation/ernie_gen_poetry/module.py +++ b/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 -- GitLab