pytorch upsample层到onnx,以及到tensorRT的转换

Posted hypnus-ly

tags:

篇首语:本文由小常识网(cha138.com)小编为大家整理,主要介绍了pytorch upsample层到onnx,以及到tensorRT的转换相关的知识,希望对你有一定的参考价值。

之前的博客介绍了upsample层转换到tensorRT出错的解决方法,就是回退onnx版本到1.5.0。虽然暂时解决了问题,但无法使用高版本的pytorch和onnx,https://www.cnblogs.com/hypnus-ly/p/12932110.html

最近又上github,发现更简单的解决方法,就是修改下upsample层初始化时的参数就可以了。 突然发现之前的博客好low啊,还是需要多多学习,发现问题的根源,找到本质解决方法。

if self.training:
    x = F.interpolate(x, scale_factor=2, mode="nearest") 
else: # hack in order to generate a simpler onnx 
    x = F.interpolate(x, size=[int(2 * x.shape[2]), int(2 * x.shape[3])], mode=nearest) # 如要转换tensorRT模型的话,需要使用 size 参数,静态的指定图片大小

 upsample函数修改方式和上述代码一样

在使用upsample函数时,经常会先定义upsample module

m = nn.Upsample(scale_factor=2, mode="nearest")

然后在前向推理时,才动态的输入数据,此时没法预先设置 size 大小,解决方法:

在进行前向推理时,先获取输入的大小,再改变函数参数:

n,c,h,w = input.size()
size = (h*2, w*2)
m.scale_factor=None
m.size = (h * 2, w * 2)
output = m(input)

注意:scale_factor 参数 和 size 参数不能同时存在,所以要将 scale_factor 设置为 None

我的环境:

pytorch-1.5, onnx-1.7, tensorRT-7.0, cuda-10.0

参考链接:

https://github.com/onnx/onnx-tensorrt/issues/457

https://github.com/onnx/onnx-tensorrt/issues/361

https://github.com/pytorch/pytorch/issues/27376 

以上是关于pytorch upsample层到onnx,以及到tensorRT的转换的主要内容,如果未能解决你的问题,请参考以下文章

pytorch C++部署模型 ONNX tensorRT

对象检测模型 (PyTorch) 到 ONNX:ONNX 推理的空输出

模型转换 PyTorch转ONNX 入门

Pytorch模型(.pth)转onnx模型(.onnx)

来自 PyTorch 模型的 ONNX 对象,无需导出

PyTorch模型导出到ONNX文件示例(LeNet-5)