pytorch->onnx->trt 踩坑记录

ERROR: ModelImporter.cpp:296 In function importModel: [5] Assertion failed: tensors.count(input_name)

问题原因:
pytorch版本为1.4.0,pytorch版本过高引起的onnx解析问题.(据悉这个解析问题会发生在trt5.0-6.0,trt7.0不会出现,详见trt6转torch1.2以上版本的onnx)

问题解决:

  1. 上面的链接中,可以直接重新编译onnx2trt,只需要载入build好的engine的同学可以直接去重新编译onnx2trt,然后用这个转onnx为trt引擎。
  2. 降torch的版本到1.1.0

[E] [TRT] Network must have at least one output

问题原因:
网络的输出为两个,转onnx的时候要明确这一点
问题解决:

error改之前:(我训练出来的网络应该有两个输出,以下代码设置的是torch.onnx.export的参数)

input_names=["input"]
output_names=["output_mask"]

改之后

input_names=["input"]
output_names=["output_conv10,""output_mask"]

将torch.onnx.export的output_names参数的输出name个数,设置的和网络模型实际输出个数相同后,就好了。

Only tuples, lists and Variables supported as JIT inputs, but got dict

问题原因:
pytorch下降到1.1.0,不支持字典的输出。
问题解决:
即原来的网络我是这么输出的


    def forward(self, x):
		...
        x5 = self.layer5(x4)
        return {"seg":x5,"mask":x4}

要改成这样


    def forward(self, x):
		...
        x5 = self.layer5(x4)
        
        return x4, x5

然后,把调用的后处理改一下,就可以了。

你可能感兴趣的:(TensorRT)