项目场景:
使用TensorRT转换ONNX模型时报错
def build_engine():
with trt.Builder(TRT_LOGGER) as builder, builder.create_network() as network, trt.OnnxParser(network, TRT_LOGGER) as parser:
builder.max_workspace_size = 1 << 30
with open(ONNX_MODEL, "rb") as model:
parser.parse(model.read())
return builder.build_cuda_engine(network)
Pytorch==1.8.0
TensorRT==7.2.3.4
问题描述:
[TensorRT] ERROR: Network must have at least one output
[TensorRT] ERROR: Network validation failed.
原因分析:
https://forums.developer.nvidia.com/t/tensorrt-error-network-must-have-at-least-one-output-tensorrt-error-network-validation-failed/156296/3
这个issue中有提到报错的原因是TensorRT的OnnxParser
解析模型失败。
解决方案:
百度之后发现,#328中@cwentland0提到,在TensorRT-7.0中,OnnxParser
只支持full-dimensions mode
,也就是需要固定Batch Size
并传一个explicitBatch
flag到create_network
方法中。
官方文档的Builder中有提到具体的实现方法:
pass a value of 1 << NetworkDefinitionCreationFlag.EXPLICIT_BATCH to create_network()。
不过具体操作中trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH
并不是一个int值,因此需要进行强制类型转换:
1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
转换之后可以成功解析模型