mirrors/mmdeploy
1
0
Fork 0
mirror of https://github.com/open-mmlab/mmdeploy.git synced 2025-01-14 08:09:43 +08:00
Code Issues Projects Releases Wiki Activity

enable TRT parse ONNX model from file (#1735)

Browse Source
This commit is contained in:
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) parser = trt.OnnxParser(network, logger)
if isinstance(onnx_model, str): 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 = '' error_msgs = ''
for error in range(parser.num_errors): for error in range(parser.num_errors):
error_msgs += f'{parser.get_error(error)}\n' error_msgs += f'{parser.get_error(error)}\n'