From 5815fbc9faab8e9eff1c7004954ed698f38c5050 Mon Sep 17 00:00:00 2001
From: wangguanzhong <jerrywgz@126.com>
Date: Thu, 24 Mar 2022 14:12:32 +0800
Subject: [PATCH] add message for export_onnx (#5444)

---
 deploy/python/infer.py       | 4 ++++
 ppdet/engine/export_utils.py | 6 ++++++
 2 files changed, 10 insertions(+)

diff --git a/deploy/python/infer.py b/deploy/python/infer.py
index 7ce83c25a..ee635c718 100644
--- a/deploy/python/infer.py
+++ b/deploy/python/infer.py
@@ -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):
diff --git a/ppdet/engine/export_utils.py b/ppdet/engine/export_utils.py
index 08871ed3a..8f55f0293 100644
--- a/ppdet/engine/export_utils.py
+++ b/ppdet/engine/export_utils.py
@@ -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':
-- 
GitLab