未验证 提交 5815fbc9 编写于 作者: W wangguanzhong 提交者: GitHub

add message for export_onnx (#5444)

上级 8db03fb3
......@@ -553,6 +553,10 @@ class PredictConfig():
self.nms = yml_conf['NMS']
if 'fpn_stride' in yml_conf:
self.fpn_stride = yml_conf['fpn_stride']
if self.arch == 'RCNN' and yml_conf.get('export_onnx', False):
print(
'The RCNN export model is used for ONNX and it only supports batch_size = 1'
)
self.print_config()
def check_model(self, yml_conf):
......
......@@ -126,7 +126,13 @@ def _dump_infer_config(config, path, image_shape, model):
'metric': config['metric'],
'use_dynamic_shape': use_dynamic_shape
})
export_onnx = config.get('export_onnx', False)
infer_arch = config['architecture']
if 'RCNN' in infer_arch and export_onnx:
logger.warning(
"Exporting RCNN model to ONNX only support batch_size = 1")
infer_cfg['export_onnx'] = True
if infer_arch in MOT_ARCH:
if infer_arch == 'DeepSORT':
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册