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

           

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