前言
构建onnx方式通常有两种:
1、通过代码转换成onnx结构,比如pytorch —> onnx
2、通过onnx 自定义结点,图,生成onnx结构
本文主要是简单学习和使用两种不同onnx结构,
下面以reshape
结点进行分析
方式
方法一:pytorch --> onnx
固定shape
import torchclass JustReshape(torch.nn.Module):def __init__(self):super(JustReshape, self).__init__()def forward(self, x):# x = x.view((x.shape[3], x.shape[1], x.shape[0], x.shape[2]))x= x.reshape(x.shape[3], x.shape[1], x.shape[0], x.shape[2])return x net = JustReshape()
model_name = 'just_reshape.onnx'#保存ONNX的文件名字
dummy_input = torch.randn(1, 31, 42, 5)
torch.onnx.export(net, dummy_input, model_name, input_names=['input'], output_names=['output'])
结果如图所示: