Fix the bug of compile model to tensorrt

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

Hide file tree

Showing 3 changed files with 4 additions and 4 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/caffe_model_file_to_onnx_model.py b/model_compiler/src/model_compiler/compilers/caffe_model_file_to_onnx_model.py
@@ -4,8 +4,8 @@
 import os
 from typing import Any, Mapping, NamedTuple, Optional, Sequence, List
 import numpy as np
-import caffe2.python.onnx.frontend
 from caffe2.proto import caffe2_pb2
+import onnx_caffe2
 from . import repository
 from .. import utilities
 from ..models.data_format import DataFormat
@@ -57,7 +57,7 @@ def compile_source(source: CaffeModelFile, config: Config) -> OnnxModel:
         input_shape.insert(0, config.max_batch_size)
         value_info[config.input_names[i]] = (config.input_type, input_shape)
 
-    onnx_model = caffe2.python.onnx.frontend.caffe2_net_to_onnx_model(predict_net, init_net, value_info)
+    onnx_model = onnx_caffe2.frontend.caffe2_net_to_onnx_model(predict_net, init_net, value_info)
 
     graph = onnx_model.graph  # pylint: disable=no-member
     return OnnxModel(model_proto=onnx_model,

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,