尧图建网站 尧图建网站 YAOTU WEB BUILD 免费咨询
ARTICLE DETAIL

资讯详情

深耕网站建设与建站编程的一线实战洞察。

为什么 .pt 能转 onnx?onnx 到底是什么?为什么部署都要先过一道 ONNX?—— 亲手跑一遍,快 5 倍的原因全在这

为什么 .pt 能转 onnx?onnx 到底是什么?为什么部署都要先过一道 ONNX?—— 亲手跑一遍,快 5 倍的原因全在这 动手跑过才敢写。本文所有结论都来自我在自己环境里亲手跑出来的真实输出一个推演、一个拍脑袋的结论都没有。环境torch 2.10.0cu128/onnx 1.21.0/onnxruntime 1.26.0。引言三个每天都想问的问题做推理部署的人几乎天天和.pt、.onnx打交道。但大多数人是会复制命令不懂原理为什么torch.onnx.export一下.pt就能变成.onnx它不是两种完全不同的文件吗ONNX 到底是什么东西一个文件格式为什么市面上的模型都要先转成 ONNX 这个中间态再转成 TensorRT / OpenVINO / RKNN为什么都说 ONNX 推理比.pt快它哪来的底气这篇我用一个 8.7 万参数的小 CNN 当实验对象把.pt和.onnx都拆开给你看再实测推理速度用数据回答这四个问题。第一章、撕开外壳.pt 和 .onnx 分别是什么格式和上一篇拆.pt一样先看文件头魔数。保存一个state_dict再用二进制读前 20 字节withopen(pt_onnx_blog/tinycnn.pt,rb)asf:print(repr(f.read(20)))我环境里真实输出bPK\x03\x04\x00\x00\x08\x08\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00PK\x03\x04是 ZIP 的魔数。所以.pt本质是个ZIP 压缩包。用zipfile看里面装了什么tinycnn/data.pkl tinycnn/.format_version tinycnn/.storage_alignment tinycnn/byteorder tinycnn/data/0 tinycnn/data/1 tinycnn/data/2 tinycnn/data/3data.pkl是 Pickle 序列化的对象结构data/0、data/1… 是每个张量的二进制数据。也就是说.pt ZIP 外壳 Pickle 内核 独立存储的张量。它存的是一堆带名字的 numpy 张量外加怎么把它们装回类的说明。那.onnx呢同样读文件头b\x08\n\x12\x07pyto...这是ProtobufProtocol Buffers的编码。\x08\n表示ir_version10\x12\x07pyto...是 producer 字段内容是pytorch_export…。也就是说ONNX 是一个用 Protobuf 序列化的、描述计算图的标准文件。用onnx库把它解析成人类能看懂的样子%input[FLOAT, 1x3x112x112] # 输入 initializers( # 权重initializer 直接挂在图里 %conv1.weight[FLOAT, 16x3x3x3] %fc.weight[FLOAT, 10x25088] ... ) %getitem Conv(...)(%input, %conv1.weight, %conv1.bias) %relu Relu(%getitem) %max_pool2d MaxPool(%relu) ...到这里两个文件的本质就清楚了.pt.onnx本质ZIP PickleProtobuf 序列化的计算图存什么张量 “如何装回类”算子节点 权重 数据流有向图依赖什么必须有对应的 Python 类定义不依赖任何框架纯数据描述计量单位张量图的节点算子一个存的是参数和类一个存的是一张算子组成的数据流图 参数。这就是两者最根本的区别也是后面所有问题的钥匙。第二章、为什么 .pt 能转成 onnx因为任何一个 PyTorch 模型本质上都长着一张计算图只是平时被藏起来了。训模型时你写的是forward()里面是一串算子Conv、ReLU、BatchNorm、MaxPool、Linear……PyTorch 底层ATen本来就把它们编排成了一张动态计算图。转 ONNX 要做的就是把这团藏着图的 Python 对象用一批真实输入喂它跑一遍把这条路径上的算子、形状、权重全部固化下来再翻译成 ONNX 的图格式。一句话torch.onnx.export做的事 用真实输入跑一遍前向trace 把算子翻译成 ONNX 标准算子 把权重固化成 initializer。我导出的图上torch.onnx在翻译时顺手做了一件事把BatchNorm 直接熔进了 Conv。看算子分布就露馅了算子节点总数: 8 算子类型分布: {Conv: 2, Relu: 2, MaxPool: 2, Reshape: 1, Gemm: 1}我的模型明明有 2 个 BatchNorm 层但导出的 ONNX 图里一个 BatchNorm 节点都没有。因为 BN 在推理时只是一组 scale/bias/shift可以折叠进前面的 Conv 权重里。这就是图优化算子融合的第一步也是 ONNX 能变快的第一个原因后面还会展开。所以为什么能转同一个模型PyTorch 眼里是类 参数ONNX 眼里是算子图 参数。导出只是做了个翻译把前者翻译成后者信息不丢还顺带优化了一下。完整导出代码可复现importtorch,torch.nnasnnclassTinyCNN(nn.Module):def__init__(self):super().__init__()self.conv1nn.Conv2d(3,16,3,padding1)self.bn1nn.BatchNorm2d(16)self.relunn.ReLU()self.poolnn.MaxPool2d(2)self.conv2nn.Conv2d(16,32,3,padding1)self.bn2nn.BatchNorm2d(32)self.fcnn.Linear(32*28*28,10)defforward(self,x):xself.pool(self.relu(self.bn1(self.conv1(x))))xself.pool(self.relu(self.bn2(self.conv2(x))))xx.view(x.size(0),-1)returnself.fc(x)modelTinyCNN().eval()xtorch.randn(1,3,112,112)torch.save(model.state_dict(),tinycnn.pt)# 存权重torch.onnx.export(model,x,tinycnn.onnx,input_names[input],output_names[output],opset_version17,)print(参数量:,sum(p.numel()forpinmodel.parameters()))# 87114注意torch.onnx.export需要你传入一个真实输入x这正是trace的要求——它得跑一遍才知道图长什么样。第三章、为什么市面上的模型都要一个 ONNX 中间态因为ONNX 不是某种推理引擎的私有格式而是深度学习模型的通用语言 / 中间表示IR。现实里训完一个模型你想部署到NVIDIA 显卡 → 用TensorRTTRT 引擎Intel CPU / 集显 → 用OpenVINOIR瑞芯微 Rockchip 板子 → 用RKNN手机/树莓派 → 用 TFLite / NCNN / MNN……你不可能每个硬件都让 PyTorch 原生支持。于是 ONNX 就成了万能中转站PyTorch(.pt) ──导出──▶ ONNX ──转换──▶ TensorRT / OpenVINO / RKNN / NCNN ... TensorFlow(.pb) ──▶ ONNX ─────────▶ 同一个 ONNX喂给任意引擎为什么需要一个中间态而不是每个引擎各自读 .pt跨框架ONNX 是开放标准PyTorch、TensorFlow、Paddle 都能导出 ONNX各家的模型都能汇到同一条部署流水线上。跨硬件每个硬件厂商只要实现读 ONNX → 转成自己的格式就一次性支持了所有框架的模型不用为每个框架各写一遍转换器。解耦训练与部署ONNX 不依赖任何 Python 类。训练端升级框架只要导出的 ONNX 标准不变部署端完全无感。所以ONNX 中间态的价值类比一下就是大家都在用通用格式比如 PDF而不是每家浏览器各出个私有格式。ONNX 就是模型界的 PDF。注意一个容易混淆的点ONNX 本身不是最终部署格式而是分发/交换格式。真正上线跑得飞快的是各引擎把 ONNX 再转出来的原生引擎文件TRT 引擎、RKNN 模型等。ONNX 站在源头和终点之间。第四章、为什么 ONNX 推理比 .pt 快—— 实测 5.11 倍先说破一个误区快的不是ONNX 这个文件而是运行它的推理引擎 ONNX RuntimeORT。.onnx文件本身只是一张图是一段描述不是编译好的二进制。速度来自用谁去执行这张图。我拿同一个模型同一张1×3×112×112输入在 CPU 上跑 200 次取平均先 warmupPyTorch eager 平均: 0.877 ms ONNX Runtime 平均: 0.172 ms 加速比: 5.11x输出一致性也验证了浮点误差可忽略最大绝对误差: 1.45e-07一个 eager 模式一个 ORT 引擎为什么差 5 倍三个原因其中第二点是核心。原因一没有 Python 解释器开销PyTorch eager 模式每执行一个算子都要走一遍Python 解释器 → 派发到 C/CUDA kernel的链路一层层函数调用、对象创建、GIL、动态形状检查。ORT 是纯 C 运行时把整张图 load 进去后一次调用sess.run()就能把整条图跑完算子之间的数据在内存里直接传递不需要来回越过 Python 边界。对小算子、小模型Python 的调度开销占比尤其大所以加速比越明显。模型超大、算子超大时Python 调度占比下降加速比会缩水——但省掉 Python 开销这点始终成立。原因二图级优化算子融合 常量折叠—— 实测 17.9 倍这是最硬核的一点。ORT 在加载 ONNX 时会做一系列图变换算子融合把Conv BatchNorm ReLU熔成一个ConvRelu、ConvAddRelu熔成一个FusedConv等减少内存读写次数和 kernel 启动次数。常量折叠图里能提前算的常量如某些 shape 计算在加载时就算完推理时不再算。死代码消除 / 冗余消除删掉不影响输出的节点。我用同一个 ONNX 文件一只开图优化、一只关掉实测图优化开启 (默认 ALL): 0.111 ms 图优化关闭 (DISABLE) : 1.993 ms 优化带来的加速 : 17.91x光图优化一项就带来 18 倍差距。这充分说明ONNX/ORT 的加速大头是把一张朴素算子图和融合后的高效算子图的区别而不是文件格式本身有多玄。原因三静态图 固定形状 → 内存规划与 kernel 选择ONNX 图结构固定、形状固定导出时定了1×3×112×112ORT 可以在会话初始化阶段就完成内存池规划、kernel 选择、多线程调度编排推理时零动态决策。而 PyTorch eager 每次前向都要现场决定形状、分配内存。小结为什么 onnx 快加速来源原理实测影响无 Python 解释器开销一次sess.run()跑完整图明显图级优化算子融合等ConvBNReLU 融合、常量折叠、死码消除单独 17.9×静态图 固定形状提前规划内存与 kernel明显公平起见补充一句PyTorch 也有一套torch.compile/ TorchScript 静态化方案来缩小差距但 ONNX 生态的跨硬件通用性 成熟的图优化让它成为部署事实标准。第五章、完整可复现脚本conda activate ysj310-gpu pipinstallonnx onnxruntime onnxscript python exp1_export.py# 导出 .pt 和 .onnxpython exp2_inspect_speed.py# 解剖图 推理速度对比python exp3_opt.py# 图优化开/关对比python exp4_magic.py# 看两个文件的魔数exp2_inspect_speed.py代码importonnximporttimeimporttorchimporttorch.nnasnnimportnumpyasnpimportonnxruntimeasort# ---------- 1. 解剖 ONNX 图 ----------monnx.load(pt_onnx_blog/tinycnn.onnx)print( ONNX 图结构 )print(输入:,[(i.name,[d.dim_valuefordini.type.tensor_type.shape.dim])foriinm.graph.input])print(输出:,[(o.name,[d.dim_valuefordino.type.tensor_type.shape.dim])foroinm.graph.output])print(算子节点总数:,len(m.graph.node))fromcollectionsimportCounter opsCounter(n.op_typeforninm.graph.node)print(算子类型分布:,dict(ops))# 打印前 8 个节点展示计算图print(前8个节点:)forninm.graph.node[:8]:print( ,n.op_type,in:,list(n.input),-,list(n.output))# ---------- 2. 重建 torch 模型对比推理速度 ----------classTinyCNN(nn.Module):def__init__(self):super().__init__()self.conv1nn.Conv2d(3,16,3,padding1)self.bn1nn.BatchNorm2d(16)self.relunn.ReLU()self.poolnn.MaxPool2d(2)self.conv2nn.Conv2d(16,32,3,padding1)self.bn2nn.BatchNorm2d(32)self.fcnn.Linear(32*28*28,10)defforward(self,x):xself.pool(self.relu(self.bn1(self.conv1(x))))xself.pool(self.relu(self.bn2(self.conv2(x))))xx.view(x.size(0),-1)returnself.fc(x)torch.manual_seed(0)modelTinyCNN().eval()# 为了公平给 torch 模型也做 torch.compile 之外的基础优化关掉梯度modelmodel.float()xtorch.randn(1,3,112,112)# warmupwithtorch.no_grad():for_inrange(10):model(x)torch.cuda.synchronize()iftorch.cuda.is_available()elseNonedefbench(fn,n200):t0time.perf_counter()for_inrange(n):fn()iftorch.cuda.is_available():torch.cuda.synchronize()return(time.perf_counter()-t0)/n*1000# mswithtorch.no_grad():torch_msbench(lambda:model(x),200)# ---------- onnxruntime CPU ----------sessort.InferenceSession(pt_onnx_blog/tinycnn.onnx,providers[CPUExecutionProvider])inputs{sess.get_inputs()[0].name:x.numpy()}for_inrange(10):sess.run(None,inputs)ort_msbench(lambda:sess.run(None,inputs),200)print(f\n 推理速度对比 (单张 112x112, CPU) )print(fPyTorch eager 平均:{torch_ms:.3f}ms)print(fONNX Runtime 平均:{ort_ms:.3f}ms)print(f加速比:{torch_ms/ort_ms:.2f}x)# ---------- 数字一致性 ----------withtorch.no_grad():torch_outmodel(x)ort_outsess.run(None,inputs)[0]print(f\n 输出一致性 )print(torch 输出 shape:,torch_out.shape,| ort 输出 shape:,ort_out.shape)diffnp.abs(torch_out.numpy()-ort_out)print(f最大绝对误差:{diff.max():.2e})总结一张表读懂问题答案验证方式.pt是什么ZIP 外壳 Pickle 内核 独立张量文件读魔数PK\x03\x04zipfile看data.pkl.onnx是什么Protobuf 序列化的计算图算子 权重读文件头\x08\nonnx库打印图为什么能转模型本来就有计算图导出trace 一遍 翻译算子 固化权重torch.onnx.export图里 BN 被熔进 Conv为什么需要中间态ONNX 是跨框架、跨硬件的通用 IR各引擎再转私有格式画出 PT→ONNX→TRT/RKNN 链路为什么 ONNX 更快① 无 Python 开销 ② 图级算子融合 ③ 静态图内存规划实测 5.11×单图优化 17.9×最后一句大白话.pt存的是参数和类.onnx存的是一张静态算子图 参数。能转是因为模型本质就是一张图需要中间态是因为 ONNX 是各种硬件引擎都认的通用语言变快不是因为文件换了个格式而是交给了 ONNX Runtime 这种纯 C 的、带图优化的静态推理引擎去执行。引申彩蛋如果你想让 ONNX 再快通常会再用onnxsimonnx-simplifier把图再撸一遍拓扑排序、常量折叠、消除冗余 reshape 等然后交给 TensorRT 生成引擎。那又是另一篇故事了。
返回列表