MSRResNet部署TensorRT踩坑

MSRResNet部署踩坑

• onnx转tensorRT的问题

onnx转tensorRT后，推理结果完全错误

原因：
在尝试直接将torch转换成tensorRT推理引擎的时候发现了问题原因

x = torch.ones((1, 3, 180, 320)).cuda()
trt_path = os.path.join(model_pool, f"{model_name}_{resolution}_trt_fp32.pth")


from models.network_msrresnet import MSRResNet1 as net
model = net(in_nc=n_channels, out_nc=n_channels, nc=64, nb=16, upscale=4)
model.load_state_dict(torch.load(model_path), strict=True)
model.eval()
for k, v in model.named_parameters():
    v.requires_grad = False
model = model.to(device)
model_trt = torch2trt(model, [x])
torch.save(model_trt.state_dict(), trt_path)

报错如下：

Warning: Encountered known unsupported method torch.nn.functional.pixel_shuffle
Warning: Encountered known unsupported method torch.nn.functional.pixel_shuffle
[12/11/2023-10:43:11] [TRT] [E] 3: HRconv:0:CONVOLUTION:GPU: at least 4 dimensions are required for input.

结论：TensorRT推理的时候无法转换nn.functional.pixel_shuffle的功能

解决方案：

将该步骤放在函数外或者在TensorRT中加入自编函数插件，手动实现pixel_shuffle