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的实现。