ONNX结构分析
ONNX将每一个网络的每一层或者说是每一个算子当作节点Node，再由这些Node去构建一个Graph，相当于是一个网络。最后将Graph和这ONNX模型的其他信息结合在一起，生成一个Model，也就是最终的.onnx的模型。构建一个简单的ONNX模型，实质上，只要构建好每一个node，然后将它们和输入输出超参数一起塞到Graph，最后转成Model就可以了。
graph{
node{
input: "1"
input: "2"
output: "12"
op_type: "Conv"
}
attribute{
name: "strides"
ints: 1
ints: 1
}
attribute{
name: "pads"
ints: 2
ints: 2
}
...
}
我们查看ONNX网络结构和参数（查看网址：https://netron.app/）
ONNX安装、使用
安装ONNX环境，在终端中执行以下命令，环境中需要提前准本 python3.6. 以下流程以ubunt 20.04 为例。
模型转换流程
超分辨率是一种提高图像、视频分辨率的算法，广泛用于图像处理或视频编辑。首先，让我们在PyTorch中创建一个SuperResolution 模型。该模型使用描述的高效子像素卷积层将图像的分辨率提高了一个放大因子。该模型将图像的YCbCr的Y分量作为输入，并以超分辨率输出放大的Y分量。
# Some standard imports
import io
import numpy as np
from torch import nn
import torch.utils.model_zoo as model_zoo
import torch.onnx
# Super Resolution model definition in PyTorch
import torch.nn as nn
import torch.nn.init as init
class SuperResolutionNet(nn.Module):
def __init__(self, upscale_factor, inplace=False):
super(SuperResolutionNet, self).__init__()
self.relu = nn.ReLU(inplace=inplace)
self.conv1 = nn.Conv2d(1, 64, (5, 5), (1, 1), (2, 2))
self.conv2 = nn.Conv2d(64, 64, (3, 3), (1, 1), (1, 1))
self.conv3 = nn.Conv2d(64, 32, (3, 3), (1, 1), (1, 1))
self.conv4 = nn.Conv2d(32, upscale_factor ** 2, (3, 3), (1, 1), (1, 1))
self.pixel_shuffle = nn.PixelShuffle(upscale_factor)
self._initialize_weights()
def forward(self, x):
x = self.relu(self.conv1(x))
x = self.relu(self.conv2(x))
x = self.relu(self.conv3(x))
x = self.pixel_shuffle(self.conv4(x))
return x
def _initialize_weights(self):
init.orthogonal_(self.conv1.weight, init.calculate_gain('relu'))
init.orthogonal_(self.conv2.weight, init.calculate_gain('relu'))
init.orthogonal_(self.conv3.weight, init.calculate_gain('relu'))
init.orthogonal_(self.conv4.weight)
# Create the super-resolution model by using the above model definition.
torch_model = SuperResolutionNet(upscale_factor=3)
由于本教程以演示为目的，因此采用下载预先训练好的权重。在导出模型之前调用torch_model.eval()或torch_model.train(False)将模型转换为推理模式很重要。因为dropout或batchnorm等运算符在推理和训练模式下的行为不同。
# Load pretrained model weights
model_url = 'https://s3.amazonaws.com/pytorch/test_data/export/superres_epoch100-44c6958e.pth'
batch_size = 1 # just a random number
# Initialize model with the pretrained weights
map_location = lambda storage, loc: storage
if torch.cuda.is_available():
map_location = None
torch_model.load_state_dict(model_zoo.load_url(model_url, map_location=map_location))
# set the model to inference mode
torch_model.eval()
要导出模型，我们调用该torch.onnx.export() 函数。这将执行模型，记录用于计算输出的运算符。因为export运行模型，我们需要提供一个输入张量x。只要它是正确的类型和大小，其中的值可以是随机的。请注意，除非指定为动态轴，否则所有输入维度的导出ONNX图中的输入大小将是固定的。在此示例中，我们使用batch_size 1的输入导出模型，但随后在dynamic_axes参数中将第一个维度指定为动态 torch.onnx.export() . 因此，导出的模型将接受大小为[batch_size, 1, 224, 224]的输入，其中batch_size可以是可变的。
# Input to the model
x = torch.randn(batch_size, 1, 224, 224, requires_grad=True)
torch_out = torch_model(x)
# Export the model
torch.onnx.export(torch_model, # model being run
x, # model input (or a tuple for multiple inputs)
"super_resolution.onnx", # where to save the model (can be a file or file-like object)
export_params=True, # store the trained parameter weights inside the model file
opset_version=10, # the ONNX version to export the model to
do_constant_folding=True, # whether to execute constant folding for optimization
input_names = ['input'], # the model's input names
output_names = ['output'], # the model's output names
dynamic_axes={'input' : {0 : 'batch_size'}, # variable length axes
'output' : {0 : 'batch_size'}})
在使用ONNX Runtime验证模型的输出之前，我们将使用ONNX的 API检查ONNX 模型。首先，onnx.load(“super_resolution.onnx”) 将加载保存的模型并输出 onnx.ModelProto结构（用于捆绑 ML 模型的顶级文件/容器格式）。然后，onnx.checker.check_model(onnx_model) 将验证模型的结构并确认模型具有有效的架构。ONNX 图的有效性通过检查模型的版本、图的结构以及节点及其输入和输出来验证。
import onnx
onnx_model = onnx.load("super_resolution.onnx")
onnx.checker.check_model(onnx_model)
import onnxruntime
ort_session = onnxruntime.InferenceSession("super_resolution.onnx")
def to_numpy(tensor):
return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()
# compute ONNX Runtime output prediction
ort_inputs = {ort_session.get_inputs()[0].name: to_numpy(x)}
ort_outs = ort_session.run(None, ort_inputs)
# compare ONNX Runtime and PyTorch results
np.testing.assert_allclose(to_numpy(torch_out), ort_outs[0], rtol=1e-03, atol=1e-05)
print("Exported model has been tested with ONNXRuntime, and the result looks good!")
处理前图片：
注：文章仅代表作者个人的观点，欢迎大家留言交流。