关键词：ONNX，模型部署，PyTorch，TensorFlow，Netron
内容摘要
- ONNX模型格式简介
- PyTorch pth模型导出ONNX
- TensorFlow pb模型导出ONNX
- 通过ONNX查看模型的权重
- 使用Netron对ONNX模型结构可视化
- 使用onnxsim对ONNX文件裁剪优化
ONNX模型格式简介
ONNX（Open Neural Network Exchange，开放式神经网络交换格式）是一种模型文件格式，它在模型训练和模型推理中间提供了中间桥梁，使得上游不同的训练框架都能导出ONNX格式的模型，给到下游不同的推理框架都可以读取ONNX进行部署。
这种基于ONNX的模型训练，中间件，再到模型推理的方式使得
- 1.ONNX将模型训练和推理解耦，任意上游训练框架和下游推理框架都可以组合搭配，而不需要用同一种框架既进行训练又进行推理
- 2.ONNX是通用的模型格式，不同训练框架输出的模型可以用ONNX作为桥梁进行转换，使得模型更方便迁移
- 3.ONNX部署兼容性极强，支持多种推理框架，支持CPU/GPU推理，支持跨语言推理
- 4.ONNX格式配合上类似ONNXRumtime等推理框架，相比于模型在原生环境的推理性能会有大幅的提升
PyTorch pth模型导出ONNX
PyTorch自带接口支持直接导出ONNX，以一个PyTorch训练得到的Bert微调模型为例，导出ONNX的示例代码如下
# 环境依赖
torch 1.12.1+cu113
onnx 1.9.0
import torch
import onnx
def convert_onnx(model, onnx_path):
input_ids = torch.LongTensor([list(range(0, 50))]).to(DEVICE)
attention_mask = torch.LongTensor([list(range(0, 50))]).to(DEVICE)
token_type_ids = torch.LongTensor([[0] * 50]).to(DEVICE)
torch.onnx.export(model,
# 输入参数的个数和顺序和forward一致，否则报错
(input_ids, attention_mask, token_type_ids),
onnx_path,
verbose=False,
opset_version=12,
input_names=['input_ids', 'attention_mask', 'token_type_ids'],
# 名字都可以自定义，但是推理的时候要和定义时的名字一致，否则报错
# 定义输出的name数量必须 <= forward输出的数量,多了报错，少了按照顺序截取输出
output_names=['out', 'prob'],
# 指定batch_size维度是动态，否则batch_size维度只能和样例数据dummy_input的batch_size一致
dynamic_axes={"input_ids": {0: "batch_size"},
"token_type_ids": {0: "batch_size"},
"attention_mask": {0: "batch_size"},
"co_vectors": {0: "batch_size"},
"out": {0: "batch_size"},
"prob": {0: "batch_size"}
})
onnx_model = onnx.load(onnx_path)
try:
onnx.checker.check_model(onnx_model)
except onnx.checker.ValidationError as e:
print(f'The model is invalid')
else:
print('The model is valid! {}'.format(onnx_path))
model = Model().to(DEVICE)
model.load_state_dict(torch.load(os.path.join(ROOT_PATH, "./model/pth_model/model.pth")))
convert_onnx(model.eval(), os.path.join(ROOT_PATH, "./model/pth_model/{}.onnx".format(int(time.time()))))
torch.onnx.export将模型由nn.Module对象转化为ONNX模型并写入路径，torch.onnx.export的设置参数如下
- model：torch.nn.Module，torch.jit.ScriptModule对象模型
- dummy input：构造一批输入数据，如果forward有多个输入则输入一个tuple，输入数据的个数和顺序和forward一致，数据的值可以随机构造
- onnx_path：转化完的ONNX文件存储路径
- opset_version：导出onnx时参考的onnx算子集版本
- input_names：输入的字段名称，和模型forward和dummy input的顺序和数量一致，名称可以用户自定义
- output_names：输出的字段名称，名称数量必须小于等于forward输出的数量，如果和forward输出数量不相等，按照forward输出的顺序截取输出，字段名称用户可以自定义，但是在模型推理的时候要和定义时的名字一致
- dynamic_axes：指定动态维度，默认在推理阶段，输入的维度必须和构造的dummy input一致，通过指定第0维度batch_size为动态维度使得模型支持任意批次大小的推理
在导出完成后借助onnx包onnx.checker.check_model对模型格式进行检查是否合法。
TensorFlow pb模型导出ONNX
TensorFlow导出ONNX需要额外的依赖包tf2onnx，tf2onnx可以通过命令方便地将pb文件转化为ONNX
# 环境依赖
tensorflow 1.15.0
tf2onnx 1.15.1
以一个TensorFlow构建的GAT模型为例，将pb转化为ONNX
>>> python -m tf2onnx.convert \
--saved-model 1690859205 \
--output model.onnx \
--inputs input_self:0,input_neigh_1:0,input_neigh_2:0,w_dropout_keep_prob:0,e_dropout_keep_prob:0,batch_normalization:0 \
--outputs softmax_out/probs:0
该命令需要传入模型的输入和输出节点名称，节点名称和TensorFLOW pb的节点名称保持一致，如果有多个输入使用逗号隔开，转化成功的日志如下。
2023-09-07 10:00:22,982 - INFO - Successfully converted TensorFlow model 1690859205 to ONNX
2023-09-07 10:00:22,982 - INFO - Model inputs: ['input_self:0', 'input_neigh_1:0', 'input_neigh_2:0', 'w_dropout_keep_prob:0', 'e_dropout_keep_prob:0', 'batch_normalization:0']
2023-09-07 10:00:22,983 - INFO - Model outputs: ['softmax_out/probs:0']
2023-09-07 10:00:22,983 - INFO - ONNX model is saved at model.onnx
通过ONNX查看模型的权重
ONNX是基于protobuf组织而成的模型结构，由下面几部分组成
| 类型 | 用途 |
|---|---|
| ModelProto | 定义了整个网络的模型结构 |
| GraphProto | 定义了模型的计算逻辑，包含了构成图的节点，这些节点组成了一个有向图结构 |
| NodeProto | 定义了每个OP的具体操作 |
| ValueInfoProto | 序列化的张量，用来保存weight和bias |
| TensorProto | 定义了输入输出形状信息和张量的维度信息 |
| AttributeProto | 定义了OP中的具体参数，比如Conv中的stride和kernel_size等 |
模型的权重存储在TensorProto类型的initializer下，通过onnx.numpy_helper.to_array可以在ONNX中拿到和PyTorch网络一样的模型权重。
以一个简单的全连接PyTorch模型转为ONNX为例，观察两者的权重参数是否一致
class Model(nn.Module):
def __init__(self):
super(Model, self).__init__()
self.linear = nn.Linear(16, 2)
self.softmax = nn.Softmax(dim=1)
def forward(self, x):
return self.softmax(self.linear(x))
model = Model()
通过named_parameters打印出线性层的权重和偏置
>>> for name, val in model.named_parameters():
... print(name, val.data)
linear.weight tensor([[ 0.2436, -0.2233, 0.1883, -0.2315, -0.1315, -0.0015, 0.1009, 0.1565,
-0.1828, 0.0095, -0.0107, -0.0619, 0.0099, -0.0578, 0.0112, -0.0310],
[-0.1352, 0.1521, 0.1886, -0.0317, 0.1550, 0.2188, -0.1523, 0.0016,
-0.2363, 0.1929, 0.0115, 0.0508, 0.0377, -0.1211, -0.0428, 0.1316]])
linear.bias tensor([-0.1867, -0.0074])
然后将模型转为ONNX格式
data = torch.tensor([list(range(16))]).float()
torch.onnx.export(model, data, "model.onnx", verbose=False, opset_version=12, input_names=['x'], output_names=['res'], dynamic_axes={"x": {0: "batch_size"}, "res": {0: "batch_size"}})
重新导入ONNX模型，使用onnx.numpy_helper.to_array在initializer中拿到权重
>>> from onnx.numpy_helper import to_array
>>> model2 = onnx.load("./model.onnx")
>>> for i in model2.graph.initializer:
... print(i.name, to_array(i))
linear.weight [[ 0.24359486 -0.22334516 0.18832973 -0.23149619 -0.13151914 -0.00149396
0.10089278 0.15651006 -0.18279949 0.0094898 -0.01069614 -0.06186351
0.0098637 -0.05777565 0.01121134 -0.03095323]
[-0.13515878 0.15211186 0.18859869 -0.03166929 0.15498772 0.21884191
-0.15227136 0.00157559 -0.23632333 0.1928516 0.01145819 0.05077234
0.03767726 -0.12113282 -0.04280344 0.131643 ]]
linear.bias [-0.1867499 -0.00738242]
比对之后两者的权重和偏置完全一致，本质上ONNX将各种上游的模型结构转化为protobuf格式，其中记录了模型中的节点名称，权重，图结构等信息，这些通用信息给到下游推理引擎进行推理。
使用Netron对ONNX模型结构可视化
Netron是神经网络可视化工具，Netron可以辅助用于观察ONNX的模型图结构，还是以上一节的简单线性模型为例，通过代码调用Netron可视化如下
# 环境依赖
# netron==7.1.6
>>> import netron
>>> netron.start("./model.onnx")
在图上可以清楚的检查模型的结构是否正确，每个节点的输入的shape信息，以及右侧每个节点的输入输出的名称。
使用onnxsim对ONNX文件裁剪优化
转换得到的ONNX可能存在冗余结构，在ONNX生态中可以使用onnx-simplifier工具对ONNX模型文件进行精简，它会扫描模型图结构，试图用恒定输出替换冗余运算符。
用PyTorch编写一个简单模型，里面人为的加入两个冗余结构
class Model(nn.Module):
def __init__(self):
super(Model, self).__init__()
self.linear1 = nn.Linear(128, 64)
self.sigmoid = nn.Sigmoid()
def forward(self, x):
concat = []
for i in range(4):
concat.append(self.linear1(x))
res = torch.concat(concat, dim=1)
return torch.reshape(torch.reshape(res, [-1, 1]), [-1, 64 * 4])
在这个网络中一个相同的线性层和输入x被for循环重复计算，在输出中两个reshape操作将数据形状改变又恢复回来，通过Netron可视化如下
>>> model = Model()
>>> data = torch.tensor([list(range(128))]).float()
>>> torch.onnx.export(model, data, "model.onnx", verbose=False, opset_version=12, input_names=['x'], output_names=['res'], dynamic_axes={"x": {0: "batch_size"}, "res": {0: "batch_size"}})
>>> # 可视化
>>> netron.start("./model.onnx")
通过torch.onnx.export导出的ONNX原封不动地把PyTorch逻辑翻译了一遍，在图上明显发现4次循环只要计算一次即可，两次Reshape在做无用功。
下面通过onnx-simplifier进行优化，代码如下
# 环境依赖
# onnxsim 0.4.33
>>> from onnxsim import simplify
>>> model2 = onnx.load("./model.onnx")
>>> model_simp, check = simplify(model2)
>>> onnx.save(model_simp, 'model_pruned.onnx')
>>> # 重新可视化
>>> netron.start("./model_pruned.onnx")
onnxsim自动删除了另外三次重复运算，使用同一个结果进行Concat，删除两次无意义的Reshape，优化后模型从34k变小为33k，onnxsim确实对冗余结构进行了精简和替换。
$ ls -lht
-rw-rw-r-- 1 root root 33K 9月 7 13:59 model_pruned.onnx
-rw-rw-r-- 1 root root 34K 9月 7 13:59 model.onnx
