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 

posted @ 2020-07-08 17:57  blackx  阅读(2379)  评论(0编辑  收藏  举报