ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

深度学习 - 24 ONNX 模型导出与量化

深度学习 - 24 ONNX 模型导出与量化 ONNX 模型导出与量化0. 先建立完整部署链路做 ASR 模型部署时,最容易出现的错误认知是:PyTorch → ONNX → INT8好像 ONNX 本身就是一个“更快的 PyTorch 模型”,然后导出时再顺手选择 INT8。实际上这几个概念解决的是不同问题:PyTorch Model │ │ Export ▼ ONNX Graph │ ├── Graph Optimization │ ├── FP16 / Mixed Precision │ └── INT8 / INT4 Quantization │ ▼ ONNX Runtime / TensorRT / QNN / ... │ ▼ CPU / GPU / NPU / Mobile可以把它理解成三层:第一层:模型表达方式PyTorch 负责训练和研究,ONNX 负责把模型的推理计算表达成一种标准化计算图。第二层:模型优化例如 Constant Folding、Operator Fusion、FP16、Mixed Precision、INT8。第三层:硬件执行例如 ONNX Runtime CPU/CUDA、TensorRT、Qualcomm QNN、OpenVINO、NNAPI 等 Execution Provider。ONNX Runtime 本身通过 Execution Provider 把图或者子图交给目标硬件对应的执行后端。因此:ONNX 是模型表示/交换格式,不是一个完整的推理引擎;量化是模型数值表示优化,也不是 ONNX 导出的同义词。这一区分是整个主题最重要的基础。1. 核心概念1.1 什么是 ONNXONNX,全称 Open Neural Network Exchange,是一种用于表示机器学习模型的开放标准。从模型角度看,一个 ONNX 模型本质上是:Model ├── Graph │ ├── Inputs │ ├── Nodes │ ├── Outputs │ └── Initializers │ ├── Operator Set / Opset ├── Data Types └── Metadata其中:Input:模型输入 TensorOutput:模型输出 TensorNode:计算节点,例如 MatMul、Conv、Add、Reshape、SoftmaxInitializer:模型参数,例如权重、biasAttribute:某些算子的静态参数Opset:节点采用哪一版本的 ONNX operator 语义ONNX 官方文档把它描述成一种专门描述数学计算的“模型语言”:一个计算可以被表达成由标准算子组成的计算图。例如:X │ ▼ MatMul │ ▼ Add │ ▼ GELU │ ▼ Y它并不关心你原来在 PyTorch 中是不是写成:y=self.linear(x)y=F.gelu(y)最终需要把这件事转换成标准计算图。1.2 ONNX 和.pt到底有什么区别这是面试非常容易被问到的问题。首先要注意:.pt不是一种严格意义上的“模型格式名称”。PyTorch 的.pt/.pth文件可以保存很多东西,例如:state_dict 完整 Module optimizer state scheduler state 训练 checkpoint 其他 Python 对象因此一个训练 checkpoint 可能包含:model weights optimizer epoch learning rate AMP scaler training state ...而 ONNX 通常是一个推理计算图 + 参数的序列化模型。典型区别:项目PyTorch.pt/.pthONNX主要用途训练 / PyTorch 推理跨框架推理计算图取决于保存内容显式保存参数通常有有Optimizer可以有通常没有Autograd可以参与不属于核心用途Training loop可以保存不保存Python 代码依赖常见运行时不应依赖原 Python ModuleRuntimePyTorchONNX Runtime / TensorRT / QNN 等硬件适配PyTorch backendRuntime / Execution ProviderC++ 部署可以,但依赖 LibTorch 等很常见因此:PyTorch.pt更接近“PyTorch 世界中的模型/检查点”,ONNX 更接近“面向推理后端的标准计算图”。现代 PyTorch 的 ONNX exporter 会捕获 PyTorch 模型的 Tensor 计算,经过torch.export的图表示,再转换为 ONNX graph。官方当前推荐的 exporter 是基于torch.export的路径。2. 为什么需要导出 ONNX2.1 最大问题:训练框架和部署环境通常不是同一个东西训练时你可能使用:PyTorch CUDA Python Lhotse icefall k2 custom Python preprocessing但部署环境可能是:C++ Android CPU TensorRT Qualcomm NPU OpenVINO Web 服务器推理服务如果直接把 PyTorch 模型带过去,就需要把整个 PyTorch 运行环境也带过去。而 ONNX 的设计目标之一就是:训练框架 ↓ 标准化模型表达 ↓ 不同 Runtime也就是:PyTorch ─────┐ TensorFlow ──┤ JAX ─────────┤ ▼ ONNX │ ┌──────┼──────┐ ▼ ▼ ▼ CPU GPU NPUONNX 官方明确把“训练框架与生产部署环境解耦”作为其重要使用场景。2.2 ONNX 的真正价值不是“文件变小”这是另一个常见误区。导出 ONNX 主要解决的是:模型表示标准化 + Runtime 解耦 + 后端优化入口。至于:模型大小降低推理更快GPU/NPU 加速INT8FP16这些是后续的优化和执行问题。所以:PyTorch → ONNX并不天然意味着:PyTorch slower ONNX faster实际性能取决于:模型结构 + ONNX graph + graph optimization + Runtime + Execution Provider + kernel implementation + Tensor layout + memory copy + target hardwareONNX Runtime 会对图做优化,并根据 Execution Provider 对图进行分区和执行。3. PyTorch → ONNX 到底发生了什么这是 Export 部分最值得掌握的内容。3.1 不是简单地“把 Python 代码保存下来”假设:classEncoder(nn.Module):defforward(self,x):x=self.linear(x)x=self.relu(x)returnxPyTorch 中你看到的是 Python Module。但 ONNX 不理解:self.linear self.relu Pythonobject它需要的是标准 Tensor computation:Input │ ▼ MatMul │ ▼ Add │ ▼ Relu │ ▼ Output因此 Export 的本质是:把 PyTorch 中可执行的 Tensor 程序捕获成一个静态/符号化的计算图,然后把这个图转换到 ONNX Operator Set。现代 PyTorch 2.6+ 的 ONNX exporter 使用torch.export路径,可以先得到一个 AOT 的 Tensor graph;这个中间图会进行 operator normalization、消除大部分 Python control flow/data structure,并记录 Shape Constraints,然后再转换为 ONNX。可以概括成:PyTorch Module │ ▼ torch.export / ExportedProgram │ │ normalized Tensor operations ▼ ONNX Graph │ ▼ Serialization │ ▼ model.onnx这比简单理解成“跑一次 forward,然后录下来”更准确。4. Example Input 为什么如此重要Export 必须知道模型的输入是什么。例如:x=torch.randn(1,100,80)对于 ASR 来说:B = 1 T = 100 F = 80那么模型可能经历:Input [B, T, F] ↓ Encoder [B, T/4, D] ↓ Transformer / Conformer [B, T/4, D] ↓ Projection [B, T/4, V]Exporter 使用 example input 来确定 Tensor 的计算关系和 Shape 信息。如果没有显式声明动态维度,那么给进去的 Shape 默认可能被当成具体 Shape。当前 PyTorch exporter 文档明确说明:如果没有指定 dynamic shapes,example input 的 Shape 会被记录成导出的输入/输出 Shape;要支持动态 Shape,需要显式声明。因此:torch.randn(1,100,80)不等于“以后自然支持任意 T”。这就是 Static Shape 和 Dynamic Shape 的来源。5. Static Shape5.1 定义Static Shape 指模型某个 Tensor 的维度在导出模型中被固定。例如:Input: [1, 100, 80]代表:B = 1 T = 100 F = 80都是固定的。5.2 为什么 Static Shape 有价值很多后端在知道 Shape 后,可以提前做更多工作:Shape known ↓ Memory planning Kernel selection Constant folding Buffer allocation Graph optimization例如:[1, 100, 80]比:[B, T, 80]更容易做静态优化。某些设备后端还会根据固定 Shape 选择更加具体的 kernel 和内存布局。ONNX Runtime 的文档也明确指出,部分 Execution Provider 在输入 Shape 已知时可以获得更好的初始化与优化机会。6. Dynamic Shape6.1 为什么需要 Dynamic ShapeASR 最大的问题之一就是:语音长度是变化的。假设有:audio A → 2 sec → T=200 audio B → 5 sec → T=500 audio C → 10 sec → T=1000如果 ONNX 模型固定:[1, 500, 80]那么无法直接接受:[1, 200, 80] [1, 1000, 80]这时候就需要动态维度。例如:[B, T, 80]其中:B = dynamic T = dynamic 80 = static这意味着:Tensor 的 rank 不变,但某些维度的实际值由 runtime 决定。7. Dynamic Shape 的底层本质Dynamic Shape 并不是:“什么 Shape 都可以。”它实际上是:在图中保留符号维度,并记录这些维度之间必须满足的约束。例如:Input [B, T, 80] ↓ Conv / Subsampling [B, T/4, 512]Exporter 不只是知道:B = ? T = ?还可能知道:output_T ≈ T / 4或者其他 Shape relation。现代 PyTorch 的 dynamic shape 机制基于 symbolic shape 和 shape constraints。7.1 Static / Dynamic 可以从 ONNX Shape 表示理解ONNX 中一个维度可以是:dim_value表示已知具体值。也可以是:dim_param表示符号维度。例如:[1, 100, 80]类似:dim_value = 1 dim_value = 100 dim_value = 80而:[B, T, 80]则类似:dim_param = B dim_param = T dim_value = 80ONNX 官方 Shape Inference 文档明确区分了固定dim_value、符号dim_param和完全未知维度。8. PyTorch 现代 Dynamic Shape 导出现代 PyTorch 2.6+、特别是新 exporter 路径中,更推荐:dynamic_shapes而不是旧 exporter 常见的:dynamic_axes当前 PyTorch 文档明确说明:dynamic_shapes用于dynamo=Truedynamic_axes主要属于旧的dynamo=Falseexporterdynamic_axes已被标记为 deprecated,现代 exporter 更推荐dynamic_shapes。典型写法:importtorchfromtorch.exportimportDim model=model.eval()x=torch.randn(1,100,80)batch=Dim("batch")time=Dim("time")torch.onnx.export(model,(x,),"encoder.onnx",input_names=["x"],output_names=["y"],dynamo=True,dynamic_shapes={"x":{0:batch,1:time,}},)这里表达的是:x : [batch, time, 80]而不是:x : [1, 100, 80]9. ASR 中到底应该 Static 还是 Dynamic这不是“Dynamic 一定更好”。要看模型的使用方式。Offline ASR例如:整句音频 ↓ Fbank ↓ Encoder ↓ Decoder音频长度不固定:[1, 250, 80] [1, 500, 80] [1, 1200, 80]那么通常需要:T = dynamicBatch 是否 dynamic,要看实际部署。Streaming ASR你的 Streaming Encoder 更典型:每次输入固定 chunk 例如: [1, 45, 80]然后:chunk ↓ encoder ↓ state下一次继续:[1, 45, 80]这时候:B = 1 T = 45 F = 80全部固定反而非常合理。也就是说:Streaming ASR 不一定应该使用 Dynamic Shape。这点面试非常容易被追问。对于固定 chunk 的 Streaming Encoder:Static Shape通常更有利于:memory planning kernel optimization predictable latency device compilation而且你的 C++ 推理接口通常本来就是:float* input ↓ [1, chunk, feature_dim]固定 Shape 可以让部署逻辑简单很多。10. Shape Inference10.1 Shape Inference 是什么假设:X = [1, 100, 80] Reshape ↓ Y = [1, 100, 512]如果模型内部存在:MatMul Reshape Transpose Concat ConvRuntime 或 ONNX 工具可能需要知道:每个 Tensor 的 ShapeShape Inference 的作用就是:根据算子的输入 Shape 和算子语义,尽可能推导中间 Tensor 的 Shape。例如:
返回列表