Fix the bug of compile model to tensorrt

Closes #491 Signed-off-by: zhangkaili <zhang.kaili@zte.com.cn>
Adlik · Jul 19, 2021 · 93b852f · 93b852f
1 parent 3436030
commit 93b852f
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/model_compiler/setup.py b/model_compiler/setup.py
@@ -58,7 +58,7 @@ def run(self):
         'tensorflow==2.4.0',
         'torch==1.7.1',
         'onnx-tf',
-        'onnx-caffe2==1.0.0',
+        'onnx-caffe2',
         'paddlepaddle',
         'paddle2onnx',
         'tensorflow_addons',

diff --git a/model_compiler/src/model_compiler/compilers/paddle_model_file_to_onnx_model.py b/model_compiler/src/model_compiler/compilers/paddle_model_file_to_onnx_model.py
@@ -5,7 +5,6 @@
 from tempfile import NamedTemporaryFile
 import onnx
 import onnx.utils
-from paddle2onnx.command import program2onnx
 from . import repository
 from .. import utilities
 
@@ -43,6 +42,7 @@ def from_env(env: Mapping[str, str]) -> 'Config':
 
 @repository.REPOSITORY.register(source_type=PaddlePaddleModelFile, target_type=OnnxModel, config_type=Config)
 def compile_source(source: PaddlePaddleModelFile, config: Config) -> OnnxModel:
+    from paddle2onnx.command import program2onnx  # pylint: disable=import-outside-toplevel
     with NamedTemporaryFile(suffix='.onnx') as onnx_file:
         program2onnx(model_dir=source.model_path,
                      save_file=onnx_file.name,