MSRResNet部署TensorRT踩坑

最新推荐文章于 2025-03-17 10:43:41 发布

qq_36549601

最新推荐文章于 2025-03-17 10:43:41 发布

阅读量243

点赞数

文章标签： pytorch 深度学习人工智能

本文链接：https://blog.csdn.net/qq_36549601/article/details/134924465

版权

文章讲述了在将MSRResNet模型从ONNX转换到TensorRT时遇到的问题，主要原因是TensorRT不支持torch.nn.functional.pixel_shuffle。解决方法是将该操作移出torch2trt转换过程，或自定义TensorRT插件实现pixel_shuffle功能。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

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