解决Pytorch AvgPool2d层不能转TensorRT的问题

问题起因

今天在导出onnx的时候，有一个AvgPool2d层是如下定义的

self.resnetv1d_down3_pool0_fwd = nn.AvgPool2d(kernel_size=(2, 2), stride=(2, 2), padding=(0, 0), count_include_pad=False, ceil_mode=True)

而它导出的onnx图中参数为：

Float(1, 80, 20, 36) = onnx::AveragePool[kernel_shape=[2, 2], pads=[0, 0, 1, 1], strides=[2, 2]]

起初也没在意pads这个参数的具体意义，直到转TensorRT出现问题。

在生成TensorRT模型时，tensorrt日志系统有如下报错，这就说明输入的pading参数有误：

the padding size is larger than or equal to the filter size for exclusive-counting pooling

定位问题

同时pooling后接的conv层也报错说输入维度不能为空，于是回头查看了onnx-tensorrt-v5.1中AvgPool的解析代码，发现这个这层输出tensor的nbDims为-1，难怪conv出错。

发现了问题所在自然就是解决它，TensorRT对onnx模型AvgPool层参数解析是这样的

  nvinfer1::DimsHW kernel_size(1, 1), strides(1, 1), beg_padding(0, 0), end_padding(0, 0);
  nvinfer1::PaddingMode paddingMode;
  get_kernel_params(node, get_DimsHW_from_CHW(dims),
                    &kernel_size, &strides, &beg_padding, &end_padding, paddingMode);

看到beg_padding和end_padding参数就大概猜到问题出在哪了，这两个参数更像来源自tensorflow的pad层，代表着第几个dimension需要被padding，可能onnx在定义avgpooling的时候就是参考的tf吧。加断点查看end_padding,果然跟onnx图里一样。由此可见错误就是因为在onnx模型中,这个avagepool的输入tensor在第2维和第3维被错误升维造成的。

解决思路

1. 用opset=10重新导出正确参数的onnx模型;
2. 修改tensorrt对onnx模型avgpool2d层的解析代码;