【pytorch】——Assertion failed: axis ＞= 0 && axis ＜ nbDims问题

硬件版本

• pytorch1.1.0
• tensorrt5.1

问题一：
Failed to export an ONNX attribute, since it’s not constant, please try to make things (e.g., kernel size) static if possible

定位：

width_output = x[i].shape[-1]
height_output= x[i].shape[-2]

y = y + F.interpolate(
    self.fuse_layers[i][j](x[j]),
    size=[height_output, width_output],
    mode='bilinear', align_corners=ALIGN_CORNERS)

这里的width_output, height_output不是常量造成的

解决方案：

• 用常数代替
• 转换成python number

shape = torch.tensor(x[i].shape).tolist()
height_output, width_output = shape[-2:]

问题二：
onnx转tensorrt时报错：
While parsing node number 550 [Gather]:
3
ERROR: /home/data/CM/8_instance_segmentation/devices/Pspnet_5.1/onnx-tensorrt/onnx2trt_utils.hpp:399 In function convert_axis:
[8] Assertion failed: axis >= 0 && axis < nbDims
ERROR: failed to parser onnx file

分析：
由问题可以看出是因为tensorrt在解析onnx时，onnx中有Gather的操作才报错的。

pytorch转onnx时的打印信息

  %1475 : Long() = onnx::Gather[axis=0](%1474, %1473), scope: HighResolutionNet/SpatialGather_Module[ocr_gather_head]
  %1476 : Long() = onnx::Constant[value={1}](), scope: HighResolutionNet/SpatialGather_Module[ocr_gather_head]
  %1477 : Tensor = onnx::Shape(%outputy1), scope: HighResolutionNet/SpatialGather_Module[ocr_gather_head]
  %1478 : Long() = onnx::Gather[axis=0](%1477, %1476), scope: HighResolutionNet/SpatialGather_Module[ocr_gather_head]
  %1480 : Tensor = onnx::Unsqueeze[axes=[0]](%1475)
  %1481 : Tensor = onnx::Unsqueeze[axes=[0]](%1478)
  %1483 : Tensor = onnx::Concat[axis=0](%1480, %1481, %1538)
  %1484 : Float(1, 17, 19200) = onnx::Reshape(%outputy1, %1483), scope: HighResolutionNet/SpatialGather_Module[ocr_gather_head]

对应到pytorch中的源码为：

batch_size, c, h, w = probs.size(0), probs.size(1), probs.size(2), probs.size(3)

出现gather的操作主要是pytorch的动态图特征，需要自动获取变量的尺寸维度。

解决方案：

• 转换成python的变量

batch_size, c = torch.tensor(probs.shape[:2]).tolist()