pytorch 转 tensorRT 踩的几个小坑_tensorrt engine set up failed

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 工具：

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')

建议打开 --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 即可，成功。