【pytorch】——使用upsample时转onnx的注意事项

背景
采用pytorch训练模型时，需要转换成其他模型时，比如：ncnn，mnn，tengine，tensorrt。。。等时需要先转onnx，然后在转换。这里容易出问题是因为pytorch是动态架构，而前面说的框架，基本都是基于静态框架，所以这时候容易出一些问题。

1. 采用pytorch动态语言的特性

pytorch模型定义

x = self.cbr(conv5)
        x = self.conv_last(x)

        if self.use_softmax:  # is True during inference
            ## pytorch
            # x = nn.functional.interpolate(
            #     x, size=segSize, mode='bilinear', align_corners=False)
            # x = nn.functional.softmax(x, dim=1)

            ## onnx 输入多大，返回多大，不能用多尺度测试
            x = nn.functional.interpolate(
                x, size=x.shape[2:], mode='nearest')
            x = nn.functional.softmax(x, dim=1)
            return x

转onnx时打印信息

  %632 : Long() = onnx::Constant[value={2}](), scope: OnnxSegmentation/C1DeepSup[decoder]
  %633 : Tensor = onnx::Shape(%631), scope: OnnxSegmentation/C1DeepSup[decoder]
  %634 : Long() = onnx::Gather[axis=0](%633, %632), scope: OnnxSegmentation/C1DeepSup[decoder] # /home/data/miniconda3/envs/torch120/lib/python3.7/site-packages/torch/nn/functional.py:2466:0
  %635 : Tensor = onnx::Constant[value={8}]()
  %636 : Tensor = onnx::Mul(%634, %635)
  %637 : Float() = onnx::Cast[to=1](%636), scope: OnnxSegmentation/C1DeepSup[decoder] # /home/data/miniconda3/envs/torch120/lib/python3.7/site-packages/torch/nn/functional.py:2466:0
  %638 : Float() = onnx::Floor(%637), scope: OnnxSegmentation/C1DeepSup[decoder] # /home/data/miniconda3/envs/torch120/lib/python3.7/site-packages/torch/nn/functional.py:2466:0
  %639 : Long() = onnx::Constant[value={3}](), scope: OnnxSegmentation/C1DeepSup[decoder]
  %640 : Tensor = onnx::Shape(%631), scope: OnnxSegmentation/C1DeepSup[decoder]
  %641 : Long() = onnx::Gather[axis=0](%640, %639), scope: OnnxSegmentation/C1DeepSup[decoder] # /home/data/miniconda3/envs/torch120/lib/python3.7/site-packages/torch/nn/functional.py:2466:0
  %642 : Tensor = onnx::Constant[value={8}]()
  %643 : Tensor = onnx::Mul(%641, %642)
  %644 : Float() = onnx::Cast[to=1](%643), scope: OnnxSegmentation/C1DeepSup[decoder] # /home/data/miniconda3/envs/torch120/lib/python3.7/site-packages/torch/nn/functional.py:2466:0
  %645 : Float() = onnx::Floor(%644), scope: OnnxSegmentation/C1DeepSup[decoder] # /home/data/miniconda3/envs/torch120/lib/python3.7/site-packages/torch/nn/functional.py:2466:0
  %646 : Tensor = onnx::Unsqueeze[axes=[0]](%638)
  %647 : Tensor = onnx::Unsqueeze[axes=[0]](%645)
  %648 : Tensor = onnx::Concat[axis=0](%646, %647)
  %649 : Tensor = onnx::Constant[value= 1  1 [ Variable[CPUFloatType]{2} ]](), scope: OnnxSegmentation/C1DeepSup[decoder]
  %650 : Tensor = onnx::Cast[to=1](%648), scope: OnnxSegmentation/C1DeepSup[decoder]
  %651 : Tensor = onnx::Shape(%631), scope: OnnxSegmentation/C1DeepSup[decoder]
  %652 : Tensor = onnx::Constant[value={0}](), scope: OnnxSegmentation/C1DeepSup[decoder]
  %653 : Tensor = onnx::Constant[value={2}](), scope: OnnxSegmentation/C1DeepSup[decoder]
  %654 : Tensor = onnx::Constant[value={4}](), scope: OnnxSegmentation/C1DeepSup[decoder]
  %655 : Tensor = onnx::Slice(%651, %653, %654, %652), scope: OnnxSegmentation/C1DeepSup[decoder]
  %656 : Tensor = onnx::Cast[to=1](%655), scope: OnnxSegmentation/C1DeepSup[decoder]
  %657 : Tensor = onnx::Div(%650, %656), scope: OnnxSegmentation/C1DeepSup[decoder]
  %658 : Tensor = onnx::Concat[axis=0](%649, %657), scope: OnnxSegmentation/C1DeepSup[decoder]
  %659 : Float(1, 16, 480, 640) = onnx::Resize[mode="nearest"](%631, %658), scope: OnnxSegmentation/C1DeepSup[decoder]

2. 采用固定输出尺寸时

x = nn.functional.interpolate( x, size=(480, 640), mode='nearest')
x = nn.functional.softmax(x, dim=1)

转onnx的打印信息

  %632 : Tensor = onnx::Constant[value= 1  1  8  8 [ Variable[CPUType]{4} ]](), scope: OnnxSegmentation/C1DeepSup[decoder]
  %633 : Float(1, 16, 480, 640) = onnx::Upsample[mode="nearest"](%631, %632), scope: OnnxSegmentation/C1DeepSup[decoder]

注意:

• 这里将upsample转成了onnx中的resize操作，这要看使用onnx模型的框架是否支持
• 其次，这和pytorch的版本是有一定关系的。以上使用pytorch1.2.0，如果用其他版本也会不一样，例如onnx-tensorrt项目，7.1支持onnx::resize，5.1只支持onnx::upsample
• 采用第二种固定尺寸时，就不会像第一种方式有那么多的辅助操作，gather，mul等

总结

• tensorrt5.1 ——pytorch1.1.0——upsampe采用固定尺寸——opset_version设置为9
• tensorrt5.1官方内置只支持 upsample nearest

有用的参考

• https://github.com/pytorch/pytorch/commit/11845cf2440c1abb0cb117bcee9532b26573e9c9