Pytorch转TensorRT中的坑

  Pytorch转ONNX再转TensorRT，其中遇到一些转的时候会出现的层需要修改的问题，这里对修改的层做一些总结。

Pytorch与TensorRT版本

  TensorRT的ONNX解释器是针对Pytorch版本编译的，如果版本不对应可能导致转模型时出现错误，如：

While parsing node number 0 [Conv]:
ERROR: ModelImporter.cpp:288 In function importModel:
[5] Assertion failed: tensors.count(input_name)
[E] failed to parse onnx file
[E] Engine could not be created

  目前已知的对应关系为：

Pytorch	TensorRT
1.2	6.0.1.5
1.3	7.0.0.11

reshape

  Pytorch中会有很多需要reshape的操作，可用的算子有：

• reshape
• view
• flatten

  前两个都是需要指定reshape后完整的尺寸大小，因此使用中需要先获取输入数据的维度，这个在Pytorch框架下使用没有问题，但用到TensorRT中就是个不定的值无法进行优化，此时会出现：

While parsing node number xxx [Gather]:
ERROR: onnx2trt_utils.hpp:277 In function convert_axis:
[8] Assertion failed: axis >= 0 && axis < nbDims

  这个在TensorRT转换中比较常见，现在主要由Pytorch将自有框架中相关操作换成flatten。需要注意：

• 如果reshape后是二维[B, C x H x W]，则使用flatten(1)；
• 如果reshape后是三维[B, C, H x W]，则使用flatten(2)；

normalize

  Pytorch中可以使用torch.nn.functional.normalize(X)进行normalize操作，但是这个函数应该是实验性质，TensorRT不支持，但还有分步进行normlized的操作，可替换为：

X = X.div(X.norm(p=2, dim=1, keepdim=True))

MaxPool3d

  MaxPool3d在TensorRT6之前是不支持的，但现在已经开始支持该操作，但有的网络定义时并不是用来对五维数据[B, C, D, H, W]进行操作，而是通过Pytorch处理数据时的bug，对四维数据[B, C, H, W]的后三维进行MaxPool，因此对这样的操作在TensorRT中以及更早的ONNX中都会报输入维度错误：

RuntimeError: [ONNXRuntimeError] : 1 : GENERAL ERROR : failed:
[ShapeInferenceError] Attribute strides has incorrect size

all concat input tensors must have the same dimensions except on the concatenation axis (0), but dimensions mismatched at input 2 at index 1. Input 0 shape: [xxx], Input 2 shape: [xxx]
While parsing node number xx [Conv]:
ERROR: builtin_op_importers.cpp:788 In function importConv:
[8] Assertion failed: (nbSpatialDims == 2 && kernel_weights.shape.nbDims == 4) || (nbSpatialDims == 3 && kernel_weights.shape.nbDims == 5)

• 如果是对后三维的数据进行三维MaxPool，可能需要进行MaxPool2d后对通道进行操作，这个还待研究；
• 如果MaxPool3d的kernel_size为(1, x, y)，stride为(1, m, n)，则MaxPool3d等价于MaxPool2d；
• 如果MaxPool3d的kernel_size为(1, x, y)，stride为(L, m, n)，则实际为对数据的第二维通道[:, C, :, :]进行筛选操作后，进行MaxPool2d.

  对于第三种情况，需要对Pytorch的Tensor进行筛选，再进行MaxPool2d，而Pytorch的通道操作目前查到的有：

• torch.unbind(tensor, dim=0)：去除某个维度
• torch.index_select(input, dim, index, out=None)：选出一维度的一些slice组合成新的tensor。指定维度的大小与index大小一致

  经测试，torch.unbind还未了解其用途的有效性，torch.index_select是有效果的，因此可用：

x = nn.MaxPool2d(kernel_size, stride=stride)(x.index_select(1, torch.tensor(range(0, self.stop, self.step))))

# 实际使用中，如果使用预训练模型，直接添加一个index_select操作会导致后续序号无法对应
# 因此需要将MaxPool3d操作替换为一个封装了MaxPool2d和index_select的操作
class MaxPool2dStep(nn.Module):
    def __init__(self, kernel_size, stride, step, stop):
        super(MaxPool2dStep, self).__init__()
        self.step = step
        self.stop = stop
        self.block = nn.Sequential(
            nn.MaxPool2d(kernel_size, stride=stride)
        )

    def forward(self, x):
        return self.block(x.index_select(1, torch.tensor(range(0, self.stop, self.step))))

mean (ReduceMean)

  torch.mean(x)可以用来获取整个Tensor的均值，但在转TensorRT时会报：

While parsing node number xx [ReduceMean]:
ERROR: onnx2trt_utils.hpp:347 In function convert_axis:
[8] Assertion failed: axis >= 0 && axis < nbDims

  尝试了几种修改方式：

• torch.mean(x, 1, keepdim=True). 能转换，但结果和未修改前不一致，因为这种均值是对第2维求均值的矩阵，而非一个常数；
• torch.mean(f_pow).expand(1, 1, f_pow.shape[2], f_pow.shape[3]). python结果相同，但无法转TensorRT，因此上种方法不是因为数位相同原因；
• torch.mean(f_pow.flatten(1), 1, keepdim=True). 能转换，且结果一致，均值是一个shape为[1, 1]的数据。

  因此，猜测TensorRT在处理除法时存在矩阵除以常数的操作，但batchsize维度（即第1维，最外层维度）和内部数据处理机制不同，因此这个常数需要至少保留2个维度，后续再进一步验证猜想。

BatchNorm1d

  对于Pytorch1.2.0转换模型为ONNX后，在序列化为TensorRT时出现：

Integer division by zero

  经过对模型进行调整，确定主要的问题出在BatchNorm1d，对于ArcFace模型的最后一层一般是是FC+BN，因此需要先将Tensor打平。但转成的ONNX后对模型进行可视化，发现最后一个BatchNorm1d其实是ONNX中的BatchNormalization操作，因此需要先将Tensor升维后再降维。

  但“Unsqueeze”这样的操作在转TensorRT时会出现问题：
Error importing onnx::Unsqueeze #180
Error importing onnx::Unsqueeze layer
  一种方法是如TensorRT对ONNX"Batchnorm1d"的替代方案，将最后的FC通过Conv实现，这样数据就不用打平，在使用BatchNorm操作完后再打平。这就要求在训练前调整好模型结构。

  另一种方法是对Batchnorm1d的操作进行重新定义。如果是已经训练好后的模型不想修改结构重新训练，就需要修改Batchnorm1d的实现方式避免转ONNX时进行升降维的操作。结合pytorch nn.BatchNorm1d 与手动python实现不一样–解决办法、class MyBatchNorm2d(nn.BatchNorm2d)，实现Batchnorm1d自定义：

class MyBatchNorm1d(nn.BatchNorm1d):
    def __init__(self, num_features):
        super(MyBatchNorm1d, self).__init__(num_features)
        self.eps = 1e-5
        self.affine = True

    def forward(self, input):
        self._check_input_dim(input)
        # calculate running estimates
        input = (input - self.running_mean) / (torch.sqrt(self.running_var + self.eps))
        if self.affine:
            input = input * self.weight + self.bias
        return input