Support mlu in infer.py

Author: qili93

deploy/python/infer.py

@@ -86,6 +86,7 @@ def auto_tune(args, imgs, img_nums):
 
 
 class Predictor:
+
  def __init__(self, args):
  """
  Prepare for prediction.
@@ -101,6 +102,8 @@ def __init__(self, args):
  self._init_cpu_config()
  elif args.device == 'npu':
  self.pred_cfg.enable_custom_device('npu')
+ elif args.device == 'mlu':
+ self.pred_cfg.enable_custom_device('mlu')
  elif args.device == 'xpu':
  self.pred_cfg.enable_xpu()
  else:
@@ -119,21 +122,22 @@ def __init__(self, args):
  if hasattr(args, 'benchmark') and args.benchmark:
  import auto_log
  pid = os.getpid()
- self.autolog = auto_log.AutoLogger(
- model_name=args.model_name,
- model_precision=args.precision,
- batch_size=args.batch_size,
- data_shape="dynamic",
- save_path=None,
- inference_config=self.pred_cfg,
- pids=pid,
- process_name=None,
- gpu_ids=0,
- time_keys=[
- 'preprocess_time', 'inference_time', 'postprocess_time'
- ],
- warmup=0,
- logger=logger)
+ self.autolog = auto_log.AutoLogger(model_name=args.model_name,
+ model_precision=args.precision,
+ batch_size=args.batch_size,
+ data_shape="dynamic",
+ save_path=None,
+ inference_config=self.pred_cfg,
+ pids=pid,
+ process_name=None,
+ gpu_ids=0,
+ time_keys=[
+ 'preprocess_time',
+ 'inference_time',
+ 'postprocess_time'
+ ],
+ warmup=0,
+ logger=logger)
 
  def _init_base_config(self):
  self.pred_cfg = PredictConfig(self.cfg.model, self.cfg.params)
@@ -224,9 +228,8 @@ def run(self, imgs_path):
  if args.benchmark:
  self.autolog.times.start()
 
- data = np.array([
- self._preprocess(p) for p in imgs_path[i:i + args.batch_size]
- ])
+ data = np.array(
+ [self._preprocess(p) for p in imgs_path[i:i + args.batch_size]])
  input_handle.reshape(data.shape)
  input_handle.copy_from_cpu(data)
 
@@ -268,35 +271,32 @@ def _save_imgs(self, results, imgs_path):
 
 def parse_args():
  parser = argparse.ArgumentParser(description='Test')
- parser.add_argument(
- "--config",
- dest="cfg",
- help="The config file.",
- default=None,
- type=str,
- required=True)
+ parser.add_argument("--config",
+ dest="cfg",
+ help="The config file.",
+ default=None,
+ type=str,
+ required=True)
  parser.add_argument(
  '--image_path',
  dest='image_path',
  help='The directory or path or file list of the images to be predicted.',
  type=str,
  default=None,
  required=True)
- parser.add_argument(
- '--batch_size',
- dest='batch_size',
- help='Mini batch size of one gpu or cpu.',
- type=int,
- default=1)
- parser.add_argument(
- '--save_dir',
- dest='save_dir',
- help='The directory for saving the predict result.',
- type=str,
- default='./output')
+ parser.add_argument('--batch_size',
+ dest='batch_size',
+ help='Mini batch size of one gpu or cpu.',
+ type=int,
+ default=1)
+ parser.add_argument('--save_dir',
+ dest='save_dir',
+ help='The directory for saving the predict result.',
+ type=str,
+ default='./output')
  parser.add_argument(
  '--device',
- choices=['cpu', 'gpu', 'xpu', 'npu'],
+ choices=['cpu', 'gpu', 'xpu', 'npu', 'mlu'],
  default="gpu",
  help="Select which device to inference, defaults to gpu.")
 
@@ -306,48 +306,45 @@ def parse_args():
  type=eval,
  choices=[True, False],
  help='Whether to use Nvidia TensorRT to accelerate prediction.')
- parser.add_argument(
- "--precision",
- default="fp32",
- type=str,
- choices=["fp32", "fp16", "int8"],
- help='The tensorrt precision.')
- parser.add_argument(
- '--min_subgraph_size',
- default=3,
- type=int,
- help='The min subgraph size in tensorrt prediction.')
+ parser.add_argument("--precision",
+ default="fp32",
+ type=str,
+ choices=["fp32", "fp16", "int8"],
+ help='The tensorrt precision.')
+ parser.add_argument('--min_subgraph_size',
+ default=3,
+ type=int,
+ help='The min subgraph size in tensorrt prediction.')
  parser.add_argument(
  '--enable_auto_tune',
  default=False,
  type=eval,
  choices=[True, False],
- help='Whether to enable tuned dynamic shape. We uses some images to collect '
+ help=
+ 'Whether to enable tuned dynamic shape. We uses some images to collect '
  'the dynamic shape for trt sub graph, which avoids setting dynamic shape manually.'
  )
- parser.add_argument(
- '--auto_tuned_shape_file',
- type=str,
- default="auto_tune_tmp.pbtxt",
- help='The temp file to save tuned dynamic shape.')
-
- parser.add_argument(
- '--cpu_threads',
- default=10,
- type=int,
- help='Number of threads to predict when using cpu.')
- parser.add_argument(
- '--enable_mkldnn',
- default=False,
- type=eval,
- choices=[True, False],
- help='Enable to use mkldnn to speed up when using cpu.')
+ parser.add_argument('--auto_tuned_shape_file',
+ type=str,
+ default="auto_tune_tmp.pbtxt",
+ help='The temp file to save tuned dynamic shape.')
+
+ parser.add_argument('--cpu_threads',
+ default=10,
+ type=int,
+ help='Number of threads to predict when using cpu.')
+ parser.add_argument('--enable_mkldnn',
+ default=False,
+ type=eval,
+ choices=[True, False],
+ help='Enable to use mkldnn to speed up when using cpu.')
 
  parser.add_argument(
  "--benchmark",
  type=eval,
  default=False,
- help="Whether to log some information about environment, model, configuration and performance."
+ help=
+ "Whether to log some information about environment, model, configuration and performance."
  )
  parser.add_argument(
  "--model_name",
@@ -356,17 +353,15 @@ def parse_args():
  help='When `--benchmark` is True, the specified model name is displayed.'
  )
 
- parser.add_argument(
- '--with_argmax',
- dest='with_argmax',
- help='Perform argmax operation on the predict result.',
- action='store_true')
- parser.add_argument(
- '--print_detail',
- default=True,
- type=eval,
- choices=[True, False],
- help='Print GLOG information of Paddle Inference.')
+ parser.add_argument('--with_argmax',
+ dest='with_argmax',
+ help='Perform argmax operation on the predict result.',
+ action='store_true')
+ parser.add_argument('--print_detail',
+ default=True,
+ type=eval,
+ choices=[True, False],
+ help='Print GLOG information of Paddle Inference.')
 
  return parser.parse_args()