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