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

最新推荐文章于 2024-09-12 17:07:23 发布

Cc不爱吃洋葱

最新推荐文章于 2024-09-12 17:07:23 发布

阅读量4.1k

点赞数 18

文章标签：人工智能 python 开发语言 transformer 学习语言模型自然语言处理

本文链接：https://blog.csdn.net/2401_85328934/article/details/139699168

版权

内容摘要

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

如何系统的去学习大模型LLM ？

作为一名热心肠的互联网老兵，我意识到有很多经验和知识值得分享给大家，也可以通过我们的能力和经验解答大家在人工智能学习中的很多困惑，所以在工作繁忙的情况下还是坚持各种整理和分享。

但苦于知识传播途径有限，很多互联网行业朋友无法获得正确的资料得到学习提升，故此将并将重要的 AI大模型资料 包括AI大模型入门学习思维导图、精品AI大模型学习书籍手册、视频教程、实战学习等录播视频免费分享出来。

😝有需要的小伙伴，可以V扫描下方二维码免费领取🆓

一、全套AGI大模型学习路线

AI大模型时代的学习之旅：从基础到前沿，掌握人工智能的核心技能！

二、640套AI大模型报告合集

这套包含640份报告的合集，涵盖了AI大模型的理论研究、技术实现、行业应用等多个方面。无论您是科研人员、工程师，还是对AI大模型感兴趣的爱好者，这套报告合集都将为您提供宝贵的信息和启示。

三、AI大模型经典PDF籍

随着人工智能技术的飞速发展，AI大模型已经成为了当今科技领域的一大热点。这些大型预训练模型，如GPT-3、BERT、XLNet等，以其强大的语言理解和生成能力，正在改变我们对人工智能的认识。那以下这些PDF籍就是非常不错的学习资源。

在这里插入图片描述

四、AI大模型商业化落地方案

阶段1：AI大模型时代的基础理解

目标：了解AI大模型的基本概念、发展历程和核心原理。
内容：
- L1.1 人工智能简述与大模型起源
- L1.2 大模型与通用人工智能
- L1.3 GPT模型的发展历程
- L1.4 模型工程
- L1.4.1 知识大模型
- L1.4.2 生产大模型
- L1.4.3 模型工程方法论
- L1.4.4 模型工程实践
- L1.5 GPT应用案例

阶段2：AI大模型API应用开发工程

目标：掌握AI大模型API的使用和开发，以及相关的编程技能。
内容：
- L2.1 API接口
- L2.1.1 OpenAI API接口
- L2.1.2 Python接口接入
- L2.1.3 BOT工具类框架
- L2.1.4 代码示例
- L2.2 Prompt框架
- L2.2.1 什么是Prompt
- L2.2.2 Prompt框架应用现状
- L2.2.3 基于GPTAS的Prompt框架
- L2.2.4 Prompt框架与Thought
- L2.2.5 Prompt框架与提示词
- L2.3 流水线工程
- L2.3.1 流水线工程的概念
- L2.3.2 流水线工程的优点
- L2.3.3 流水线工程的应用
- L2.4 总结与展望

阶段3：AI大模型应用架构实践

目标：深入理解AI大模型的应用架构，并能够进行私有化部署。
内容：
- L3.1 Agent模型框架
- L3.1.1 Agent模型框架的设计理念
- L3.1.2 Agent模型框架的核心组件
- L3.1.3 Agent模型框架的实现细节
- L3.2 MetaGPT
- L3.2.1 MetaGPT的基本概念
- L3.2.2 MetaGPT的工作原理
- L3.2.3 MetaGPT的应用场景
- L3.3 ChatGLM
- L3.3.1 ChatGLM的特点
- L3.3.2 ChatGLM的开发环境
- L3.3.3 ChatGLM的使用示例
- L3.4 LLAMA
- L3.4.1 LLAMA的特点
- L3.4.2 LLAMA的开发环境
- L3.4.3 LLAMA的使用示例
- L3.5 其他大模型介绍