[Fix] fix the error of build model in export_onnx.py (PaddlePaddle#2910)

PeiyuLau · Jan 11, 2023 · 6ec7417 · 6ec7417
1 parent 09c1306
commit 6ec7417
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/tools/model/export_onnx.py b/tools/model/export_onnx.py
@@ -36,7 +36,7 @@
 import onnx
 import onnxruntime
 
-from paddleseg.cvlibs import Config
+from paddleseg.cvlibs import Config, SegBuilder
 from paddleseg.utils import logger, utils
 
 
@@ -80,8 +80,11 @@ def check_and_run_onnx(onnx_model_path, input_data):
 
 
 def export_onnx(args):
+    assert args.config is not None, \
+        'Please set --config path/to/yml'
     cfg = Config(args.config)
-    model = cfg.model
+    builder = SegBuilder(cfg)
+    model = builder.model
     if args.model_path is not None:
         utils.load_entire_model(model, args.model_path)
         logger.info('Loaded trained params of model successfully')