pytorch转tensorRT踩的几个小坑

1、版本不匹配

[E] [TRT] Layer:Where_51's output can not be used as shape tensor.
[E] [TRT] Network validation failed.
[E] Engine creation failed.
[E] Engine set up failed.

这实际是由于pytorch与TensorRT版本不匹配，我的TensorRT是7.0，pytorch应该是1.4，但我用了1.7

因此需要用1.7重新读取权重文件，然后用老的方式保存，再用onnx导出

def main():
    input_shape = (3, 416, 416)
    model_onnx_path = "yolov4tiny.onnx"

    # model = torch.hub.load('mateuszbuda/brain-segmentation-pytorch', 'unet',
    #                        in_channels=3, out_channels=1, init_features=32, pretrained=True)
    model = YoloBody(3, 12).cuda()
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    dummy_input = torch.randn(1, 3, 416, 416, device=device)
    # 用1.7版本读取权重
    state_dict = torch.load('logs/Epoch120-Total_Loss0.5324-Val_Loss0.8735.pth', map_location=device)
    model.load_state_dict(state_dict)
    # 保存成1.4版本支持的格式
    torch.save(model.state_dict(), 'logs/for_onnx.pth', _use_new_zipfile_serialization=False)
    
    # Python解释器换成torch1.4的环境，重新读取，导出pytorch1.4版本对应的onnx
    state_dict = torch.load('logs/for_onnx.pth', map_location=device)
    model.load_state_dict(state_dict)
    model.train(False)

    inputs = ['input_1']
    outputs = ['output_1', 'output_2']
    dynamic_axes = {'input_1': {0: 'batch'}, 'output_1': {0: 'batch'}}
    torch.onnx.export(model,
                      dummy_input,
                      model_onnx_path,
                      export_params=True,
                      opset_version=11,
                      do_constant_folding=True,
                      input_names=inputs, output_names=outputs,
                      dynamic_axes=None)

这样操作产生的onnx文件才能被TensorRT转换

使用TensorRT的OSS工具：

trtexec --explicitBatch --onnx=yolov4tiny.onnx --saveEngine=yolov4tiny.trt --fp16 --workspace=10240 --verbose

建议打开--verbose，转换过程会很慢，verbose打印日志看着能安心点，要不然盯着屏幕会以为卡死了，发慌

 

2、上采样scale问题

[5] Assertion failed: ctx->tensors().count(inputName)

YOLO的上采样阶段，Pytorch使用opset=11的onnx会导致upsample层里增加一个constant节点，所以TensorFlowRT转换失败，期间参考pytorch经onnx使用TensorRT部署转换踩坑记录中提到的方法，无效

尝试多个版本的Pytorch与onnx后，upsample层的问题依然解决不了，最后参考https://github.com/Tianxiaomo/pytorch-YOLOv4这个实现，在inference时不使用torch自己的插值函数，而是自己重写，成功导出TensorRT

class Upsample(nn.Module):
    def __init__(self):
        super(Upsample, self).__init__()

    def forward(self, x, target_size, inference=False):
        assert (x.data.dim() == 4)
        # _, _, tH, tW = target_size

        if inference:

            #B = x.data.size(0)
            #C = x.data.size(1)
            #H = x.data.size(2)
            #W = x.data.size(3)

            return x.view(x.size(0), x.size(1), x.size(2), 1, x.size(3), 1).\
                    expand(x.size(0), x.size(1), x.size(2), target_size[2] // x.size(2), x.size(3), target_size[3] // x.size(3)).\
                    contiguous().view(x.size(0), x.size(1), target_size[2], target_size[3])
        else:
            return F.interpolate(x, size=(target_size[2], target_size[3]), mode='nearest')

3、数据类型错误

Unsupported ONNX data type: DOUBLE (2)

由于上面的改动中，出现了除法，导致TRT识别双精度，下图中出现的cast节点即为问题。很神奇的是用上面原始的代码不会出现这个问题，换了我自己的模型就有问题了

根据https://github.com/onnx/onnx-tensorrt/issues/400#issuecomment-730240546的说法

尝试了下并没什么用。

为了解决这个问题，直接把target_size[3] // x.size(3)换成结果2即可，成功。