4.1 构建onnx结构模型

模型"/>

4.1 构建onnx结构模型

前言

构建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'])

结果如图所示：