mirrors
/
mmdeploy
mirror of https://github.com/open-mmlab/mmdeploy.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

enable TRT parse ONNX model from file (#1735)

pull/1740/head
AllentDan 2023-02-09 15:54:37 +08:00 committed by GitHub
parent f5a05b5219
commit a3311b0bbb
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 6 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
mmdeploy/backend/tensorrt/utils.py
View File

@ -164,9 +164,13 @@ def from_onnx(onnx_model: Union[str, onnx.ModelProto],
parser = trt.OnnxParser(network, logger)
if isinstance(onnx_model, str):
onnx_model = onnx.load(onnx_model)
parse_valid = parser.parse_from_file(onnx_model)
elif isinstance(onnx_model, onnx.ModelProto):
parse_valid = parser.parse(onnx_model.SerializeToString())
else:
raise TypeError('Unsupported onnx model type!')
if not parser.parse(onnx_model.SerializeToString()):
if not parse_valid:
error_msgs = ''
for error in range(parser.num_errors):
error_msgs += f'{parser.get_error(error)}\n'