模型部署系列：ONNX模型格式的构建、分析和优化指南

作者: xiaogp | 来源:发表于2023-09-10 18:34 被阅读0次

关键词：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模型中间件示意图

这种基于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")

netron对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

模型部署系列：ONNX模型格式的构建、分析和优化指南

内容摘要

ONNX模型格式简介

PyTorch pth模型导出ONNX

TensorFlow pb模型导出ONNX

通过ONNX查看模型的权重

使用Netron对ONNX模型结构可视化

使用onnxsim对ONNX文件裁剪优化

相关文章

网友评论

延伸阅读

深度阅读

栏目导航

热点阅读