先简单定义一个网络结构
class Net(torch.nn.Module):
def __init__(self):
super(Net, self).__init__()
self.ClassifierLayer = nn.Sequential(
nn.Conv2d(in_channels=3, out_channels=4, kernel_size=3, stride=1),
nn.ReLU(),
nn.Conv2d(in_channels=4, out_channels=4, kernel_size=3, stride=1),
nn.MaxPool2d(kernel_size=2),
)
def forward(self, x):
y = self.ClassifierLayer(x)
y = F.interpolate(y,size=x.shape[2:])
return y
转换代码(固定输入尺寸)
#实例化网络结构
model = Net()
#定义一个模拟的输入
input_data = torch.randn([1,3,320,320])
#参数1:模型 参数2:模拟的输入数据 参数3:导出的onnx绝对路径 #参数4:是否打印导出过程中的详细信息
torch.onnx.export(model, input_data, "Net.onnx", verbose=True)
接下来可视化一下导出后的onnx网络结构
注意看右边的输入部分[1,3,320,320]
动态输入尺寸转换代码
input_name = 'input'
output_name = 'output'
torch.onnx.export(model,
input_data,
"simplenet.onnx",
opset_version=11,
input_names=[input_name],
output_names=[output_name],
dynamic_axes={
input_name: {0: 'batch_size', 2: 'in_width', 3: 'int_height'},
output_name: {0: 'batch_size', 2: 'out_width', 3: 'out_height'}}
)
同样注意右边的输入部分,现在变成了[batch_size,3,in_width,int_height]
至此最基础的pytorch转onnx就全部完成了,所有的网络结构都可套用这两段代码,不过如果遇到了转换失败算子不支持的问题,可以尝试修改opset_version参数,或者使用其他结构去代替不支持的操作
再补充一下onnxsim:
如果你想把onnx再转其他的模型,就拿ncnn来举例,上图中的Gather在转换过程中就会出现不支持转到ncnn的错误信息。
推荐使用onnxsim,不论任何模型都可以用,当然能否sim成功还要看你的模型复杂程度,不过目前遇到的大部分模型都是可以经过sim规避很多问题的。
项目地址:onnxsim
假设你已经安装好了onnxsim,那么
onnxsim xxx.onnx xxx-sim.onnx
经过onnxsim之后,你的onnx模型就会简洁很多了,以静态输入尺寸举例如下图
可以看到ncnn不支持的Gather层已经被简化掉了。
注意:onnxsim的用法这里只说明了最基本的,其余的更多用法自行查看参数含义去使用。