【记录·持续更新】Pytorch2ONNX之踩坑记录～～

1、onnx不支持多分支逻辑判断：

class Model(torch.nn.Module):
    def __init__(self, feature_dims=768):
        super().__init__()
        self.fc = torch.nn.Linear(feature_dims * 2, feature_dims)
        self.fc2 = torch.nn.Linear(feature_dims * 2, feature_dims)

    def forward(self, data, flag):
        # flag = torch.ones(1)
        output = self.fc(data)
        if not flag.item():
            return self.fc2(data)
        return output

如图，如果将上述模型转换为onnx，将会根据转换时输入的flag选择其中一条支路进行转换，最终转换得到的onnx将不再具备if的判断功能。

2、onnx需要利用到每一个输入参数：

class Model(torch.nn.Module):
    def __init__(self, feature_dims=768):
        super().__init__()
        self.fc = torch.nn.Linear(feature_dims * 2, feature_dims)

    def forward(self, data1, data2):
        output = self.fc(data1)
        return output

如图，输入参数data2没有被使用到，或者当data2被使用了但是没有相关的变量作为输出，那么转onnx时会报输入参数非法的错误。

3、转换onnx动态尺寸

x = torch.randn(1, 3, 32, 32)
y = torch.randn(1, 10)

input_tensor = (x,
                y)

torch.onnx.export(model, 
                  input_tensor, 
                  onnx_model_path, 
                  verbose=False, 
                  input_names=["input1", "input2"], 
                  output_names=["output"], 
                  opset_version=11,
                  dynamic_axes={
                          "input1": {0: "batch_size", 2: "height", 3: "width"},
                          "output": {0: "batch_size"}
                  })

如图，上述转onnx时指定两个输入的名字分别为input1,、input2；输出只有1个，名字为output，然后指定input1的第0/2/3维度为动态，不指定input2的维度（即固定维度）；指定输出output的第0维为动态尺寸。

实际测试时，输入的input1维度为(?,3,?,?)，输入的input2维度为(1, 10)，输出维度为(?, N)

？代表动态尺寸。

4、ReLU6算子

# ReLU6函数的pytorch实现：
out = nn.ReLU6(inplace=inplace)(x)
out = F.relu6(x+3, inplace=True)

# 转onnx时，需要修改写法如下：
def relu6(x, inplace=True):
    out = nn.ReLU(inplace=inplace)(x)
    out[out > 6] = 6
    return out

# 转onnx的ReLU6类实现：
class ReLU6:
    def __init__(self, inplace=True):
        super(ReLU6, self).__init__()
        self.inplace = inplace
        self.relu = nn.ReLU(inplace=inplace)

    def forward(self, x):
        out = self.relu(x)
        out[out > 6] = 6
        return out


def hardsigmoid(x):
    return F.relu6(x + 3., inplace=True) / 6.

# 转onnx的hardsigmoid的修改写法：
def hardsigmoid(x):
    return F.hardtanh(x + 3, 0., 6.) / 6.

使用pytorch的relu6转onnx时会报错，在转换时需要修改对应relu6的实现。