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升维后再降维。
Pytorch转TensorRT中的坑_第1张图片
  但“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

  
  
  
  

你可能感兴趣的:(开源架构,深度学习)