def init_predictor(args):
if args.model_dir is not "":
config = Config(args.model_dir)
else:
config = Config(args.model_file, args.params_file)
config.enable_memory_optim()
if args.tune:
config.collect_shape_range_info(shape_file)
if args.use_gpu:
config.enable_use_gpu(1000, 0)
if args.use_trt:
# using dynamic shpae mode, the max_batch_size will be ignored.
config.enable_tensorrt_engine(workspace_size=1 << 30,
max_batch_size=1,
min_subgraph_size=5,
precision_mode=PrecisionType.Float32,
use_static=False,
use_calib_mode=False)
if args.tuned_dynamic_shape:
config.enable_tuned_tensorrt_dynamic_shape(shape_file, True)
else:
# If not specific mkldnn, you can set the blas thread.
# The thread num should not be greater than the number of cores in the CPU.
config.set_cpu_math_library_num_threads(4)
config.enable_mkldnn()
predictor = create_predictor(config)
return predictor
def get_config(self, model, params, tuned=False):
config = Config()
config.set_model_buffer(model, len(model), params, len(params))
config.enable_use_gpu(100, 0)
config.set_optim_cache_dir('tuned_test')
if tuned:
config.collect_shape_range_info('shape_range.pbtxt')
else:
config.enable_tensorrt_engine(
workspace_size=1024,
max_batch_size=1,
min_subgraph_size=0,
precision_mode=paddle.inference.PrecisionType.Float32,
use_static=True,
use_calib_mode=False)
config.enable_tuned_tensorrt_dynamic_shape('shape_range.pbtxt',
True)
return config
class Predictor:
def __init__(self, args):
"""
Prepare for prediction.
The usage and docs of paddle inference, please refer to
https://paddleinference.paddlepaddle.org.cn/product_introduction/summary.html
"""
self.args = args
self.cfg = DeployConfig(args.cfg)
self._init_base_config()
if args.device == 'cpu':
self._init_cpu_config()
else:
self._init_gpu_config()
self.predictor = create_predictor(self.pred_cfg)
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)
def _init_base_config(self):
self.pred_cfg = PredictConfig(self.cfg.model, self.cfg.params)
if not self.args.print_detail:
self.pred_cfg.disable_glog_info()
self.pred_cfg.enable_memory_optim()
self.pred_cfg.switch_ir_optim(True)
def _init_cpu_config(self):
"""
Init the config for x86 cpu.
"""
logger.info("Use CPU")
self.pred_cfg.disable_gpu()
if self.args.enable_mkldnn:
logger.info("Use MKLDNN")
# cache 10 different shapes for mkldnn
self.pred_cfg.set_mkldnn_cache_capacity(10)
self.pred_cfg.enable_mkldnn()
self.pred_cfg.set_cpu_math_library_num_threads(self.args.cpu_threads)
def _init_gpu_config(self):
"""
Init the config for nvidia gpu.
"""
logger.info("Use GPU")
self.pred_cfg.enable_use_gpu(100, 0)
precision_map = {
"fp16": PrecisionType.Half,
"fp32": PrecisionType.Float32,
"int8": PrecisionType.Int8
}
precision_mode = precision_map[self.args.precision]
if self.args.use_trt:
logger.info("Use TRT")
self.pred_cfg.enable_tensorrt_engine(workspace_size=1 << 30,
max_batch_size=1,
min_subgraph_size=50,
precision_mode=precision_mode,
use_static=False,
use_calib_mode=False)
if use_auto_tune(self.args) and \
os.path.exists(self.args.auto_tuned_shape_file):
logger.info("Use auto tuned dynamic shape")
allow_build_at_runtime = True
self.pred_cfg.enable_tuned_tensorrt_dynamic_shape(
self.args.auto_tuned_shape_file, allow_build_at_runtime)
else:
logger.info("Use manual set dynamic shape")
min_input_shape = {"x": [1, 3, 100, 100]}
max_input_shape = {"x": [1, 3, 2000, 3000]}
opt_input_shape = {"x": [1, 3, 512, 1024]}
self.pred_cfg.set_trt_dynamic_shape_info(
min_input_shape, max_input_shape, opt_input_shape)
def run(self, imgs):
if not isinstance(imgs, (list, tuple)):
imgs = [imgs]
num = len(imgs)
input_names = self.predictor.get_input_names()
input_handle = self.predictor.get_input_handle(input_names[0])
output_names = self.predictor.get_output_names()
output_handle = self.predictor.get_output_handle(output_names[0])
results = []
args = self.args
if not os.path.exists(args.save_dir):
os.makedirs(args.save_dir)
for i in range(0, num, args.batch_size):
if args.benchmark:
self.autolog.times.start()
data = np.array(
[self._preprocess(img) for img in imgs[i:i + args.batch_size]])
input_handle.reshape(data.shape)
input_handle.copy_from_cpu(data)
if args.benchmark:
self.autolog.times.stamp()
self.predictor.run()
results = output_handle.copy_to_cpu()
if args.benchmark:
self.autolog.times.stamp()
results = self._postprocess(results)
if args.benchmark:
self.autolog.times.end(stamp=True)
self._save_imgs(results, imgs)
logger.info("Finish")
def _preprocess(self, img):
return self.cfg.transforms(img)[0]
def _postprocess(self, results):
if self.args.with_argmax:
results = np.argmax(results, axis=1)
return results
def _save_imgs(self, results, imgs):
for i in range(results.shape[0]):
result = get_pseudo_color_map(results[i])
basename = os.path.basename(imgs[i])
basename, _ = os.path.splitext(basename)
basename = f'{basename}.png'
result.save(os.path.join(self.args.save_dir, basename))
class Predictor:
def __init__(self, args):
"""
Prepare for prediction.
The usage and docs of paddle inference, please refer to
https://paddleinference.paddlepaddle.org.cn/product_introduction/summary.html
"""
self.args = args
self.cfg = DeployConfig(args.cfg)
self._init_base_config()
if args.device == 'cpu':
self._init_cpu_config()
else:
self._init_gpu_config()
self.predictor = create_predictor(self.pred_cfg)
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)
def _init_base_config(self):
self.pred_cfg = PredictConfig(self.cfg.model, self.cfg.params)
if not self.args.print_detail:
self.pred_cfg.disable_glog_info()
self.pred_cfg.enable_memory_optim()
self.pred_cfg.switch_ir_optim(True)
def _init_cpu_config(self):
"""
Init the config for x86 cpu.
"""
logger.info("Using CPU")
self.pred_cfg.disable_gpu()
if self.args.enable_mkldnn:
logger.info("Using MKLDNN")
# cache 1- different shapes for mkldnn
self.pred_cfg.set_mkldnn_cache_capacity(10)
self.pred_cfg.enable_mkldnn()
self.pred_cfg.set_cpu_math_library_num_threads(self.args.cpu_threads)
def _init_gpu_config(self):
"""
Init the config for nvidia gpu.
"""
logger.info("using GPU")
self.pred_cfg.enable_use_gpu(100, 0)
precision_map = {
"fp16": PrecisionType.Half,
"fp32": PrecisionType.Float32,
"int8": PrecisionType.Int8
}
precision_mode = precision_map[self.args.precision]
if self.args.use_trt:
logger.info("Use TRT")
self.pred_cfg.enable_tensorrt_engine(workspace_size=1 << 30,
max_batch_size=1,
min_subgraph_size=300,
precision_mode=precision_mode,
use_static=False,
use_calib_mode=False)
if use_auto_tune(self.args) and \
os.path.exists(self.args.auto_tuned_shape_file):
logger.info("Use auto tuned dynamic shape")
allow_build_at_runtime = True
self.pred_cfg.enable_tuned_tensorrt_dynamic_shape(
self.args.auto_tuned_shape_file, allow_build_at_runtime)
else:
logger.info("Use manual set dynamic shape")
min_input_shape = {"x": [1, 3, 100, 100]}
max_input_shape = {"x": [1, 3, 2000, 3000]}
opt_input_shape = {"x": [1, 3, 512, 1024]}
self.pred_cfg.set_trt_dynamic_shape_info(
min_input_shape, max_input_shape, opt_input_shape)
def run(self, imgs, trimaps=None, imgs_dir=None):
self.imgs_dir = imgs_dir
num = len(imgs)
input_names = self.predictor.get_input_names()
input_handle = {}
for i in range(len(input_names)):
input_handle[input_names[i]] = self.predictor.get_input_handle(
input_names[i])
output_names = self.predictor.get_output_names()
output_handle = self.predictor.get_output_handle(output_names[0])
args = self.args
for i in tqdm.tqdm(range(0, num, args.batch_size)):
# warm up
if i == 0 and args.benchmark:
for _ in range(5):
img_inputs = []
if trimaps is not None:
trimap_inputs = []
trans_info = []
for j in range(i, i + args.batch_size):
img = imgs[i]
trimap = trimaps[i] if trimaps is not None else None
data = self._preprocess(img=img, trimap=trimap)
img_inputs.append(data['img'])
if trimaps is not None:
trimap_inputs.append(
data['trimap'][np.newaxis, :, :])
trans_info.append(data['trans_info'])
img_inputs = np.array(img_inputs)
if trimaps is not None:
trimap_inputs = (
np.array(trimap_inputs)).astype('float32')
input_handle['img'].copy_from_cpu(img_inputs)
if trimaps is not None:
input_handle['trimap'].copy_from_cpu(trimap_inputs)
self.predictor.run()
results = output_handle.copy_to_cpu()
results = results.squeeze(1)
for j in range(args.batch_size):
trimap = trimap_inputs[
j] if trimaps is not None else None
result = self._postprocess(results[j],
trans_info[j],
trimap=trimap)
# inference
if args.benchmark:
self.autolog.times.start()
img_inputs = []
if trimaps is not None:
trimap_inputs = []
trans_info = []
for j in range(i, i + args.batch_size):
img = imgs[i]
trimap = trimaps[i] if trimaps is not None else None
data = self._preprocess(img=img, trimap=trimap)
img_inputs.append(data['img'])
if trimaps is not None:
trimap_inputs.append(data['trimap'][np.newaxis, :, :])
trans_info.append(data['trans_info'])
img_inputs = np.array(img_inputs)
if trimaps is not None:
trimap_inputs = (np.array(trimap_inputs)).astype('float32')
input_handle['img'].copy_from_cpu(img_inputs)
if trimaps is not None:
input_handle['trimap'].copy_from_cpu(trimap_inputs)
if args.benchmark:
self.autolog.times.stamp()
self.predictor.run()
if args.benchmark:
self.autolog.times.stamp()
results = output_handle.copy_to_cpu()
results = results.squeeze(1)
for j in range(args.batch_size):
trimap = trimap_inputs[j] if trimaps is not None else None
result = self._postprocess(results[j],
trans_info[j],
trimap=trimap)
self._save_imgs(result, imgs[i + j])
if args.benchmark:
self.autolog.times.end(stamp=True)
logger.info("Finish")
def _preprocess(self, img, trimap=None):
data = {}
data['img'] = img
if trimap is not None:
data['trimap'] = trimap
data['gt_fields'] = ['trimap']
data = self.cfg.transforms(data)
return data
def _postprocess(self, alpha, trans_info, trimap=None):
"""recover pred to origin shape"""
if trimap is not None:
trimap = trimap.squeeze(0)
alpha[trimap == 0] = 0
alpha[trimap == 255] = 1
for item in trans_info[::-1]:
if item[0] == 'resize':
h, w = item[1][0], item[1][1]
alpha = cv2.resize(alpha, (w, h),
interpolation=cv2.INTER_LINEAR)
elif item[0] == 'padding':
h, w = item[1][0], item[1][1]
alpha = alpha[:, :, 0:h, 0:w]
else:
raise Exception("Unexpected info '{}' in im_info".format(
item[0]))
return alpha
def _save_imgs(self, alpha, img_path):
ori_img = cv2.imread(img_path)
alpha = (alpha * 255).astype('uint8')
if self.imgs_dir is not None:
img_path = img_path.replace(self.imgs_dir, '')
name, ext = os.path.splitext(img_path)
if name[0] == '/':
name = name[1:]
alpha_save_path = os.path.join(args.save_dir, 'alpha/', name + '.png')
clip_save_path = os.path.join(args.save_dir, 'clip/', name + '.png')
# save alpha
mkdir(alpha_save_path)
cv2.imwrite(alpha_save_path, alpha)
# save clip image
mkdir(clip_save_path)
alpha = alpha[:, :, np.newaxis]
clip = np.concatenate([ori_img, alpha], axis=-1)
cv2.imwrite(clip_save_path, clip)
