📅 创建时间：2026-06-03 🏷️ 标签：#ONNX #TFLite #HLO #SavedModel #opset #模型转换 #FlatBuffer #ProtocolBuffer
📚 前置知识：[[05-operation-semantics]]（算子语义） 📚 相关知识：[[03-graph-representation]]（计算图表示）[[07-graph-optimization-passes]]（图优化）
模型前端格式：ONNX、TFLite、HLO 与 SavedModel / Model Frontend Formats: ONNX, TFLite, HLO, and SavedModel
┌──────────────────────────────────────────────────────────────────────────────┐
│ 🎯 场景：同一个 PyTorch ResNet-50 如何部署到 4 个不同的目标平台 │
├──────────────────────────────────────────────────────────────────────────────┤
│ │
│ 你有一个 PyTorch ResNet-50 模型，想部署到： │
│ │
│ 1️⃣ NVIDIA GPU 服务器 → 用 TensorRT │
│ 2️⃣ Android 手机 → 用 TFLite │
│ 3️⃣ Intel 服务器 → 用 OpenVINO │
│ 4️⃣ Web 浏览器 → 用 ONNX.js │
│ │
│ 每个目标平台都需要不同的模型格式。 │
│ 你开始做 PyTorch → ONNX → 各平台格式的转换， │
│ 但每个步骤都有坑：opset 版本不对、动态 Shape 丢失、 │
│ 自定义算子无法转换。到底应该怎么管理这些转换？ │
│ │
└──────────────────────────────────────────────────────────────────────────────┘第1节：模型格式全景——为什么需要这么多格式
1.1 深度学习模型格式的碎片化
深度学习领域存在大量的模型格式，每个格式都是为了解决特定问题而设计的：
python
"""
为什么需要这么多模型格式？
┌─────────────────────────────────────────────────────────────────────────────┐
│ 模型格式的演进 │
├─────────────────────────────────────────────────────────────────────────────┤
│ │
│ 早期：框架绑定格式 │
│ ═════════════════════════════════════════════ │
│ TensorFlow → SavedModel, Freeze Graph │
│ PyTorch → State Dict (.pt), ScriptModule │
│ MXNet → Symbol + Params │
│ │
│ 问题：模型只能在本框架使用，无法跨框架分享 │
│ │
│ ════════════════════════════════════════════════════════════════════ │
│ │
│ 中期：跨框架格式 │
│ ════════════════════════════════════════════════════════════════════ │
│ ONNX → 通用图表示（Facebook + Microsoft） │
│ NNEF → NVIDIA 推动的格式（已较少使用） │
│ │
│ 问题：opset 版本混乱，转换质量参差不齐 │
│ │
│ ════════════════════════════════════════════════════════════════════ │
│ │
│ 近期：部署优化格式 │
│ ════════════════════════════════════════════════════════════════════ │
│ TFLite FlatBuffer → 移动端/嵌入式 │
│ TensorRT → NVIDIA GPU │
│ OpenVINO IR → Intel CPU/GPU/VPU │
│ TVM Relay → 通用部署 │
│ IREE → MLIR 统一格式 │
│ │
└─────────────────────────────────────────────────────────────────────────────┘
"""1.2 格式选择决策树
┌─────────────────────────────────────────────────────────────────────────────┐
│ 模型格式选择决策树 │
├─────────────────────────────────────────────────────────────────────────────┤
│ │
│ 你有模型要部署？ │
│ │ │
│ ▼ │
│ ┌───────────────────────┐ │
│ │ 训练框架是什么？ │ │
│ └───────────────────────┘ │
│ │ │ │ │
│ ▼ ▼ ▼ │
│ PyTorch TensorFlow 其他 │
│ │ │ │ │
│ ▼ ▼ ▼ │
│ ┌─────────────┐ ┌───────────┐ ┌────────────┐ │
│ │ PyTorch Edge │ │TF SavedModel│ │ ONNX │ │
│ │ TFLite/ONNX │ │ TFLite │ │ (最通用) │ │
│ └─────────────┘ └───────────┘ └────────────┘ │
│ │ │ │ │
│ ▼ ▼ ▼ │
│ ┌─────────────┐ ┌───────────┐ ┌────────────┐ │
│ │ 部署目标？ │ │ 部署目标？ │ │ 部署目标？ │ │
│ └─────────────┘ └───────────┘ └────────────┘ │
│ │ │ │ │ │ │ │
│ ▼ ▼ ▼ ▼ ▼ ▼ │
│ ┌────┐ ┌────┐ ┌────┐ ┌────┐ ┌────┐ ┌────┐ │
│ │NVIDIA│ │移动│ │ Intel│ │ Web │ │NPU │ │ CPU │ │
│ │ GPU │ │设备│ │硬件 │ │浏览器│ │专用 │ │通用 │ │
│ └──┬──┘ └──┬──┘ └──┬──┘ └──┬──┘ └──┬──┘ └──┬──┘ │
│ ▼ ▼ ▼ ▼ ▼ ▼ │
│ ┌──────┐ ┌──────┐ ┌─────┐ ┌─────┐ ┌─────┐ ┌──────┐ │
│ │TensorRT│ │TFLite│ │OpenVINO│ │ONNX.js│ │TVM/IREE│ │TVM/IREE│ │
│ └──────┘ └──────┘ └─────┘ └─────┘ └─────┘ └──────┘ │
│ │
└─────────────────────────────────────────────────────────────────────────────┘第2节：XLA HLO——Google 内部格式
2.1 HLO 是什么
HLO（High-Level Optimizer） 是 Google XLA（Accelerated Linear Algebra）编译器的中间表示。
python
"""
HLO 的特点：
1. 面向批量操作：
- HLO 操作符都是批量级别的
- 如 HloModule::Create 会创建一个完整的计算模块
2. 设备无关：
- HLO 本身不绑定到特定硬件
- 同一个 HLO 可以 lower 到 CPU、GPU、TPU
3. SSA 形式：
- 所有值都是 immutable
- 新值通过操作产生
4. 布局控制：
- 显式控制数据在内存中的布局
- NCHW vs NHWC 可配置
示例 HLO Module：
HloModule matmul
ENTRY main {
a = f32[128,256] parameter(0)
b = f32[256,512] parameter(1)
ROOT c = f32[128,512] dot(a, b),
lhs_contracting_dims={1},
rhs_contracting_dims={0}
}
"""2.2 HLO vs 其他格式
| 特性 | HLO | ONNX | Linalg |
|---|---|---|---|
| 抽象层级 | 中层 | 中层 | 高层 |
| 硬件目标 | TPU, GPU, CPU | 通用 | 通用 |
| 控制流 | HLO while/conditional | 算子化 | SCF |
| 内存模型 | 显式 buffer | 隐式 | MemRef |
| 谁用 | JAX, TF XLA | 跨框架 | MLIR 生态 |
| 可读性 | 高 | 中 | 高 |
第3节：ONNX——通用的跨框架格式
3.1 ONNX 概述
ONNX（Open Neural Network Exchange） 是由 Facebook（Meta）和 Microsoft 联合推出的跨框架模型格式。
python
"""
ONNX 核心概念：
1. 计算图（Graph）：
- 由节点（Node）组成的有向无环图
- 节点之间通过 named inputs/outputs 连接
2. 算子（Operator / Opset）：
- 每个算子有规范的定义和语义
- Opset 是算子集的版本（如 opset 13, 17, 18）
- 不同版本可能添加新算子或改变语义
3. 类型系统：
- Tensor：带 shape 和 dtype
- Sequence：可变长序列
- Map：键值对
示例 ONNX Graph：
ir_version: 8
opset_import: [""]
input {
name: "x"
type { tensor_type { elem_type: FLOAT, shape { dim { dim_value: 3 } } } }
}
node {
input: ["x"]
output: ["y"]
op_type: "Relu"
}
output {
name: "y"
type { tensor_type { elem_type: FLOAT, shape { dim { dim_value: 3 } } } }
}
"""3.2 Opset 版本管理
python
# ============================================================
# ONNX Opset 版本详解
# ============================================================
print("=" * 60)
print("ONNX Opset 版本")
print("=" * 60)
"""
ONNX Opset 版本历史：
| Opset | 发布年份 | 主要新特性 |
|-------|----------|------------|
| 1-10 | 2017 | 基础算子集 |
| 11 | 2019 | Dynamic Slice, Scan |
| 12 | 2020 | Quantize, Round |
| 13 | 2021 | 改进 Cast, ReduceMin/Max |
| 14 | 2022 | 改进 Linear, Dropout |
| 15 | 2022 | 改进 If, Sequence |
| 16 | 2023 | Attention, SkipLayerNorm |
| 17 | 2023 | 改进 Split, GRU/LSTM |
| 18 | 2024 | 改进 GatherElements |
| 19 | 2025 | 最新稳定版 |
版本选择建议：
- opset 11：兼容性最好，但功能最少
- opset 13-14：平衡点，推荐用于大多数场景
- opset 17+：新模型需要新特性时使用
"""
print()
print("如何指定 Opset 版本：")
print("-" * 40)
print("""
PyTorch → ONNX：
torch.onnx.export(
model,
x,
"model.onnx",
opset_version=17 # 指定 opset 版本
)
TensorFlow → ONNX：
import tf2onnx
tf2onnx.convert.from_graph_def(
graph_def,
opset=17
)
指定太低版本的问题：
- 某些新算子不支持
- 需要用多个基础算子模拟一个复杂算子
- 导致模型变大、效率降低
指定太高版本的问题：
- 推理运行时可能不支持
- 需要升级 ONNX Runtime
""")3.3 PyTorch 导出 ONNX 实战
python
import torch
import torch.nn as nn
import torch.onnx
# ============================================================
# PyTorch → ONNX 导出完整示例
# ============================================================
class SimpleResBlock(nn.Module):
"""简单的残差块"""
def __init__(self, channels):
super().__init__()
self.conv1 = nn.Conv2d(channels, channels, 3, padding=1)
self.bn1 = nn.BatchNorm2d(channels)
self.conv2 = nn.Conv2d(channels, channels, 3, padding=1)
self.bn2 = nn.BatchNorm2d(channels)
self.relu = nn.ReLU(inplace=True)
def forward(self, x):
residual = x
out = self.relu(self.bn1(self.conv1(x)))
out = self.bn2(self.conv2(out))
out = self.relu(out + residual)
return out
class SimpleResNet(nn.Module):
"""简化的 ResNet"""
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(3, 64, 7, stride=2, padding=3)
self.bn1 = nn.BatchNorm2d(64)
self.relu = nn.ReLU(inplace=True)
self.maxpool = nn.MaxPool2d(3, stride=2, padding=1)
self.layer1 = SimpleResBlock(64)
self.layer2 = SimpleResBlock(64)
self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
self.fc = nn.Linear(64, 1000)
def forward(self, x):
x = self.relu(self.bn1(self.conv1(x)))
x = self.maxpool(x)
x = self.layer1(x)
x = self.layer2(x)
x = self.avgpool(x)
x = torch.flatten(x, 1)
x = self.fc(x)
return x
# 创建模型
model = SimpleResNet()
model.eval()
# 准备输入
x = torch.randn(1, 3, 224, 224)
# ============================================================
# 导出选项详解
# ============================================================
print("=" * 60)
print("torch.onnx.export 完整参数")
print("=" * 60)
print("""
torch.onnx.export(
model, # 要导出的模型
args, # 示例输入（用于 tracing）
f, # 输出文件路径
export_params=True, # 是否导出模型参数
opset_version=17, # ONNX opset 版本
input_names=['input'], # 输入名称列表
output_names=['output'], # 输出名称列表
dynamic_axes={ # 动态维度配置
'input': {0: 'batch_size'},
'output': {0: 'batch_size'}
},
verbose=True, # 是否打印详细信息
do_constant_folding=True # 是否做常量折叠
)
""")
# 导出（推荐配置）
torch.onnx.export(
model,
x,
"resnet18_simplified.onnx",
export_params=True,
opset_version=17,
input_names=["input"],
output_names=["output"],
dynamic_axes={
"input": {0: "batch_size"},
"output": {0: "batch_size"}
},
do_constant_folding=True,
verbose=False
)
print("✓ 模型已导出到 resnet18_simplified.onnx")3.4 ONNX 导出常见问题
python
# ============================================================
# ONNX 导出陷阱和解决方案
# ============================================================
print("=" * 60)
print("ONNX 导出常见问题")
print("=" * 60)
print("""
问题1：动态 Shape 变成静态
┌─────────────────────────────────────────────────────────────────────────────┐
│ 原因：Tracing 只执行一次，只记录执行时的 Shape │
│ │
│ 解决方案： │
│ 1. 多次 tracing（传入不同的 batch size） │
│ 2. 使用 dynamic_axes 指定哪些维度应该是动态的 │
│ 3. 使用 symbolic_trace（实验性） │
│ │
│ ```python │
│ torch.onnx.export( │
│ model, x, "model.onnx", │
│ dynamic_axes={ │
│ 'input': {0: 'batch_size'}, │
│ 'output': {0: 'batch_size'} │
│ } │
│ ) │
│ ``` │
└─────────────────────────────────────────────────────────────────────────────┘
问题2：控制流丢失
┌─────────────────────────────────────────────────────────────────────────────┐
│ 原因：Tracing 无法捕获动态控制流（如 if x.sum() > 0） │
│ │
│ 解决方案： │
│ 1. 使用 TorchScript 代替 Tracing │
│ 2. 重构代码，避免数据依赖的控制流 │
│ 3. 使用固定的 trace_input，强制走特定分支 │
│ │
│ ```python │
│ # TorchScript 方式 │
│ scripted = torch.jit.script(model) │
│ torch.onnx.export(scripted, x, "model.onnx") │
│ ``` │
└─────────────────────────────────────────────────────────────────────────────┘
问题3：自定义算子无法转换
┌─────────────────────────────────────────────────────────────────────────────┐
│ 原因：ONNX 不支持该算子，或者 exporter 不知道如何映射 │
│ │
│ 解决方案： │
│ 1. 注册自定义算子的 ONNX 映射（symbolic function） │
│ 2. 使用 onnx.helper 创建算子节点 │
│ 3. 替换为 ONNX 支持的等价格式 │
│ │
│ ```python │
│ @torch.onnx.symbolic_helper.parse_args('v', 'v', 'v', 'v') │
│ def my_custom_op(g, a, b, c): │
│ # 定义如何把 custom_op 转换为 ONNX 算子组合 │
│ return g.op('Mul', a, b) │
│ ``` │
└─────────────────────────────────────────────────────────────────────────────┘
问题4：opset 版本不兼容
┌─────────────────────────────────────────────────────────────────────────────┐
│ 原因：使用的算子在指定的 opset 版本中不存在 │
│ │
│ 解决方案： │
│ 1. 升级 opset 版本 │
│ 2. 或者替换为低版本支持的等价算子 │
│ │
│ ```python │
│ # 查看哪些算子在哪些版本可用 │
│ import onnx │
│ model = onnx.load("model.onnx") │
│ print(model.opset_import) # 列出所有 opset 版本 │
│ ``` │
└─────────────────────────────────────────────────────────────────────────────┘
""")第4节：TFLite FlatBuffer——移动端优化
4.1 FlatBuffer 简介
FlatBuffer 是 Google 开发的高效序列化格式，比 Protocol Buffer 更高效，TFLite 使用它来存储模型。
python
"""
FlatBuffer vs Protocol Buffer：
| 特性 | FlatBuffer | Protocol Buffer |
|------|------------|-----------------|
| 访问方式 | 直接访问，无需解析 | 需要解析整个结构 |
| 内存 | 只访问需要的字段 | 需要完整解析 |
| 速度 | 更快 | 较慢 |
| 文件大小 | 更小 | 较大 |
| 用途 | TFLite 模型 | TF SavedModel |
TFLite FlatBuffer 结构：
┌────────────────────────────────────────┐
│ Header (模型元信息) │
│ - 版本号 │
│ - 算子数量 │
│ - 子图数量 │
├────────────────────────────────────────┤
│ Subgraphs (子图列表) │
│ - 每个子图包含： │
│ - 输入/输出 tensors │
│ - 算子列表 │
│ - 寄存器表 │
├────────────────────────────────────────┤
│ Tensors (张量表) │
│ - 每个张量的 shape、dtype、buffer │
├────────────────────────────────────────┤
│ Buffers (数据缓冲) │
│ - 存储权重和常量数据 │
└────────────────────────────────────────┘
"""4.2 TFLite 转换实战
python
# ============================================================
# TensorFlow → TFLite 转换
# ============================================================
print("=" * 60)
print("TensorFlow → TFLite 转换")
print("=" * 60)
print("""
方法1：使用 TFLite Converter（推荐）
import tensorflow as tf
# 加载 SavedModel 或 Keras 模型
model = tf.saved_model.load('saved_model_dir')
# 或
model = tf.keras.models.load_model('model.h5')
# 转换为 TFLite
converter = tf.lite.TFLiteConverter.from_saved_model('saved_model_dir')
# 或
converter = tf.lite.TFLiteConverter.from_keras_model(model)
# 可选：优化配置
converter.optimizations = [tf.lite.Optimize.DEFAULT] # 量化优化
converter.target_spec.supported_types = [tf.float16] # float16 量化
converter.target_spec.supported_ops = [
tf.lite.OpsSet.TFLITE_BUILTINS, # TFLite 内置算子
tf.lite.OpsSet.SELECT_TF_OPS # 需要 TF ops（更大体积）
]
# 转换
tflite_model = converter.convert()
# 保存
with open('model.tflite', 'wb') as f:
f.write(tflite_model)
""")
print()
print("方法2：动态量化 vs 静态量化")
print("-" * 40)
print("""
动态量化（Dynamic Range Quantization）：
- 最简单，不需要校准数据
- 仅量化权重，激活仍是 float32
- 模型大小减少 ~4x
- 速度提升有限
静态量化（Static Range Quantization）：
- 需要校准数据
- 权重和激活都量化
- 模型大小减少 ~4x
- 速度显著提升
float16 量化：
- 不需要校准数据
- 权重转换为 float16
- 模型大小减少 ~2x
- 兼容性较好（大多数设备支持）
""")4.3 TFLite Delegate 机制
python
"""
TFLite Delegate 允许使用硬件特定加速：
┌─────────────────────────────────────────────────────────────────────────────┐
│ TFLite Delegate 架构 │
├─────────────────────────────────────────────────────────────────────────────┤
│ │
│ TFLite Runtime │
│ │ │
│ ┌────────────────────────┼────────────────────────┐ │
│ │ │ │ │
│ ▼ ▼ ▼ │
│ ┌──────────┐ ┌──────────┐ ┌──────────┐ │
│ │ NNAPI │ │ GPU │ │ Hexagon │ │
│ │ Delegate │ │ Delegate │ │ Delegate │ │
│ └──────────┘ └──────────┘ └──────────┘ │
│ │ │ │ │
│ ▼ ▼ ▼ │
│ ┌──────────┐ ┌──────────┐ ┌──────────┐ │
│ │ Android │ │ Adreno │ │ Qualcomm │ │
│ │ NPU │ │ GPU │ │ DSP │ │
│ └──────────┘ └──────────┘ └──────────┘ │
│ │
└─────────────────────────────────────────────────────────────────────────────┘
使用 Delegate：
import tflite_runtime.interpreter as tflite
# 加载模型
interpreter = tflite.Interpreter('model.tflite')
# 使用 GPU Delegate
from tflite_runtime.interpreter import load_delegate
gpu_delegate = load_delegate('libtensorflowlite_gpu_delegate.so')
interpreter = tflite.Interpreter('model.tflite',
experimental_delegates=[gpu_delegate])
# 使用 NNAPI Delegate（Android）
nnapi_delegate = load_delegate('libnnapi.so')
interpreter = tflite.Interpreter('model.tflite',
experimental_delegates=[nnapi_delegate])
"""第5节：PyTorch 导出路径详解
5.1 PyTorch 导出的多种路径
python
# ============================================================
# PyTorch 导出路径对比
# ============================================================
print("=" * 60)
print("PyTorch 导出路径对比")
print("=" * 60)
print("""
┌─────────────────────────────────────────────────────────────────────────────┐
│ PyTorch 导出路径 │
├─────────────────────────────────────────────────────────────────────────────┤
│ │
│ PyTorch Model │
│ │ │
│ ┌───────────────────┼───────────────────┐ │
│ │ │ │ │
│ ▼ ▼ ▼ │
│ ┌───────────┐ ┌───────────┐ ┌───────────┐ │
│ │ TorchScript │ │ ONNX │ │ Exported │ │
│ │ (jit) │ │ │ │ Program │ │
│ └─────┬─────┘ └─────┬─────┘ └─────┬─────┘ │
│ │ │ │ │
│ ▼ ▼ ▼ │
│ ┌───────────┐ ┌───────────┐ ┌───────────┐ │
│ │ .pt/.pth │ │ .onnx │ │ .pt2 │ │
│ │ (TorchScript) │ │ │(torch.export)│ │
│ └───────────┘ └───────────┘ └───────────┘ │
│ │
└─────────────────────────────────────────────────────────────────────────────┘
""")
print()
print("路径1：TorchScript")
print("-" * 40)
print("""
TorchScript 是 PyTorch 的静态子集，支持：
- 训练好的 nn.Module
- if/while/for 控制流
- torch 算子
导出方法：
# 方法1：tracing（只捕获一次执行路径）
traced = torch.jit.trace(model, example_input)
traced.save('model.pt')
# 方法2：script（解析 Python AST）
scripted = torch.jit.script(model)
scripted.save('model.pt')
# 方法3：script 部分模块
class PartiallyScripted(nn.Module):
def __init__(self):
super().__init__()
self.layer = nn.Linear(10, 10)
@torch.jit.script
def complex_func(self, x):
if x.sum() > 0:
return x * 2
else:
return x / 2
优缺点：
✅ 支持控制流
✅ PyTorch 原生
❌ 不能直接部署到非 PyTorch 环境
❌ 优化有限
""")
print()
print("路径2：ONNX Export")
print("-" * 40)
print("""
torch.onnx.export 是最常用的导出方式：
torch.onnx.export(model, x, 'model.onnx', ...)
优势：
✅ 跨框架支持
✅ ONNX Runtime 高效推理
✅ 丰富的工具链支持
劣势：
❌ Tracing 限制（控制流问题）
❌ opset 版本管理复杂
❌ 自定义算子需要处理
""")
print()
print("路径3：torch.export（新版，Python 3.10+）")
print("-" * 40)
print("""
torch.export 是 PyTorch 2.0 引入的新导出 API：
from torch.export import export
exported = export(model, (example_input,))
exported.save('model.pt2')
特点：
✅ 真正的静态图（使用 torch.compile）
✅ 支持动态 Shape（通过 dynamic_shapes）
✅ 更严格的验证
示例：
from torch.export import export, Dim
# 定义动态维度
batch_dim = Dim('batch', min=1, max=1024)
exported = export(
model,
(example_input,),
dynamic_shapes={'input': {0: batch_dim}}
)
exported.save('model.pt2')
""")第6节：常见转换陷阱汇总
6.1 转换工具链对比
| 工具 | 输入 | 输出 | 特点 |
|---|---|---|---|
torch.onnx.export | PyTorch | ONNX | 官方推荐 |
tf2onnx | TF | ONNX | TF → ONNX 转换 |
onnx2torch | ONNX | PyTorch | ONNX → PyTorch |
onnx-tf | ONNX | TF | ONNX → TF |
onnx2tensorrt | ONNX | TensorRT | ONNX → TRT |
onnx-simplifier | ONNX | ONNX | 简化 ONNX 模型 |
onnxruntime | ONNX | - | 推理引擎 |
6.2 动态 Shape 问题
python
# ============================================================
# 动态 Shape 问题的完整解决方案
# ============================================================
print("=" * 60)
print("动态 Shape 问题")
print("=" * 60)
print("""
问题：PyTorch 模型支持动态 batch size，但导出后变成静态
原因：
torch.onnx.export 使用 tracing，会用示例输入的 shape
如果示例是 (1, 3, 224, 224)，导出的模型只接受这个 shape
解决方案1：多次 Tracing（暴力法）
# 用多个不同的 batch size 导出
torch.onnx.export(
model,
torch.randn(1, 3, 224, 224), # batch=1
'model_batch1.onnx'
)
torch.onnx.export(
model,
torch.randn(8, 3, 224, 224), # batch=8
'model_batch8.onnx'
)
# 局限性：无法覆盖所有可能的 batch size
解决方案2：dynamic_axes（推荐）
torch.onnx.export(
model,
torch.randn(1, 3, 224, 224),
'model.onnx',
dynamic_axes={
'input': {0: 'batch_size'}, # 第0维是动态的
'output': {0: 'batch_size'} # 输出也跟着变
}
)
# 结果：batch_size 变成 '?'，运行时可以是任意值
解决方案3：symbolic_trace（实验性）
# PyTorch 2.0+ 支持
from torch._dynamo import make_boxed_func
model = make_boxed_func(model) # 先包装
# 注意：symbolic_trace 仍然有限制
""")
print()
print("序列长度动态问题：")
print("-" * 40)
print("""
NLP 模型的序列长度也是动态的：
输入: (batch, seq_len, hidden)
seq_len 可能是 32, 64, 128, 512...
解决方案：
torch.onnx.export(
model,
(torch.randn(1, 128, 768),), # 使用较长的序列作为示例
'model.onnx',
input_names=['input_ids', 'attention_mask'],
dynamic_axes={
'input_ids': {1: 'sequence_length'},
'attention_mask': {1: 'sequence_length'},
'output': {1: 'sequence_length'}
}
)
注意：
- 使用你能想到的最大序列长度作为示例
- ONNX Runtime 会为中间节点分配最大序列长度的内存
- 内存使用 = max_seq_len * batch_size * hidden_size
""")6.3 自定义算子处理
python
# ============================================================
# 自定义算子处理
# ============================================================
print()
print("=" * 60)
print("自定义算子处理")
print("=" * 60)
print("""
场景：PyTorch 模型使用了自定义算子（如自定义的 xxx算子）
问题：
ONNX 标准算子集不包含 xxx算子
torch.onnx.export 不知道如何转换
解决方案1：注册自定义 symbolic
┌─────────────────────────────────────────────────────────────────────────────┐
│ │
│ import torch.onnx │
│ │
│ @torch.onnx.symbolic_helper.parse_args('v', 'v', 'v', 'v') │
│ def custom_op_symbolic(g, input1, input2, weight, bias): │
│ # 定义如何转换为 ONNX │
│ # 方式1：使用 ONNX 内置算子组合 │
│ tmp = g.op('Add', input1, weight) │
│ tmp = g.op('Mul', tmp, input2) │
│ return g.op('Add', tmp, bias) │
│ │
│ # 注册 symbolic │
│ torch.onnx.symbolic_registry.register_opset_symbolic( │
│ 'custom_op', 'custom_op_symbolic', 'ai.onnx', 17 │
│ ) │
│ │
└─────────────────────────────────────────────────────────────────────────────┘
解决方案2：替换为等价的 ONNX 算子
┌─────────────────────────────────────────────────────────────────────────────┐
│ │
│ # 在导出前把自定义算子替换为 ONNX 支持的算子 │
│ │
│ class ModelWithCustomOp(nn.Module): │
│ def __init__(self): │
│ super().__init__() │
│ self.custom = CustomOp() │
│ │
│ def forward(self, x): │
│ return self.custom(x) # 导出前需要替换 │
│ │
│ # 替换方案： │
│ class ModelWithOnnxOp(nn.Module): │
│ def __init__(self): │
│ super().__init__() │
│ # 用等价的 PyTorch 算子替代 │
│ self.scale = nn.Parameter(torch.ones(1)) │
│ self.bias = nn.Parameter(torch.zeros(1)) │
│ │
│ def forward(self, x): │
│ return x * self.scale + self.bias # 等价于 custom_op │
│ │
└─────────────────────────────────────────────────────────────────────────────┘
解决方案3：使用 ONNX 自定义算子
┌─────────────────────────────────────────────────────────────────────────────┐
│ │
│ # 如果必须在 ONNX 中保留自定义算子 │
│ │
│ import onnx │
│ from onnx import helper, TensorProto │
│ │
│ # 定义自定义算子节点 │
│ node_def = helper.make_node( │
│ 'CustomOp', # 算子类型 │
│ inputs=['input'], # 输入 │
│ outputs=['output'], # 输出 │
│ name='custom_node', # 节点名称 │
│ domain='custom_domain', # 命名空间 │
│ attribute1=1.0, # 属性 │
│ attribute2='value' │
│ ) │
│ │
│ # 然后手动构建图 │
│ graph_def = helper.make_graph( │
│ [node_def], │
│ 'custom_graph', │
│ [helper.make_tensor_value_info('input', TensorProto.FLOAT, [3])], │
│ [helper.make_tensor_value_info('output', TensorProto.FLOAT, [3])] │
│ ) │
│ │
└─────────────────────────────────────────────────────────────────────────────┘
""")第7节：模型转换工具链实战
7.1 完整转换流程
┌─────────────────────────────────────────────────────────────────────────────┐
│ 完整模型转换流程（以 ResNet-50 为例） │
├─────────────────────────────────────────────────────────────────────────────┤
│ │
│ Step 1: PyTorch → ONNX │
│ ════════════════════════════════════════════════════════════════════════ │
│ │
│ PyTorch Model │
│ │ │
│ ▼ │
│ torch.onnx.export( │
│ model, │
│ x, │
│ 'resnet50.onnx', │
│ opset_version=17, │
│ dynamic_axes={'input': {0: 'batch'}, 'output': {0: 'batch'}} │
│ ) │
│ │
│ Step 2: ONNX 验证 │
│ ════════════════════════════════════════════════════════════════════════ │
│ │
│ resnet50.onnx │
│ │ │
│ ▼ │
│ onnx.checker.check_model(model) │
│ # 验证算子支持、shape 约束等 │
│ │
│ Step 3: ONNX 简化（可选） │
│ ════════════════════════════════════════════════════════════════════════ │
│ │
│ pip install onnx-simplifier │
│ python -m onnxsim resnet50.onnx resnet50_simplified.onnx │
│ # 移除冗余的 identity、cast 等节点 │
│ │
│ Step 4: 根据目标选择转换 │
│ ════════════════════════════════════════════════════════════════════════ │
│ │
│ ┌─────────────────────────────────────────┐ │
│ │ ONNX Model │ │
│ └─────────────────────────────────────────┘ │
│ │ │
│ ┌─────────────────┼─────────────────┬─────────────────┐ │
│ │ │ │ │ │
│ ▼ ▼ ▼ ▼ │
│ ┌────────┐ ┌────────┐ ┌────────┐ ┌────────┐ │
│ │TensorRT│ │ TFLite │ │OpenVINO│ │ ONNX.js│ │
│ │(.plan) │ │(.tflite)│ │ (.xml) │ │(.wasm) │ │
│ └────────┘ └────────┘ └────────┘ └────────┘ │
│ │
│ Step 5: 各平台具体转换命令 │
│ ════════════════════════════════════════════════════════════════════════ │
│ │
│ TensorRT: │
│ trtexec --onnx=resnet50.onnx --saveEngine=resnet50.plan │
│ │
│ TFLite: │
│ # 先转换为 TensorFlow，然后 TFLite Converter │
│ python -m tf2onnx.convert --input resnet50.onnx --output resnet50_tf │
│ # 或直接使用 paddle2onnx 的逆向 │
│ │
│ OpenVINO: │
│ mo --input_model resnet50.onnx --output_dir ./ │
│ │
│ ONNX.js: │
│ onnxruntime-web 直接加载 .onnx 文件 │
│ │
└─────────────────────────────────────────────────────────────────────────────┘7.2 TensorRT 转换示例
python
# ============================================================
# TensorRT 转换
# ============================================================
print("=" * 60)
print("TensorRT 转换")
print("=" * 60)
print("""
TensorRT 支持的输入格式：
1. ONNX (.onnx)
2. TensorRT Engine (.plan)
3. UFF (Universal Framework Format) - 已弃用
4. Caffe Model
使用 trtexec 命令行：
# 基本转换
trtexec --onnx=resnet50.onnx --saveEngine=resnet50.plan
# 指定精度
trtexec --onnx=resnet50.onnx --saveEngine=resnet50_fp16.plan \\
--fp16
# 指定 workspace 大小
trtexec --onnx=resnet50.onnx --saveEngine=resnet50.plan \\
--workspace=4096 # 4GB
# 动态 batch
trtexec --onnx=resnet50.onnx --saveEngine=resnet50.plan \\
--minShapes=input:1x3x224x224 \\
--optShapes=input:8x3x224x224 \\
--maxShapes=input:32x3x224x224
Python API：
import tensorrt as trt
logger = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(logger)
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
parser = trt.OnnxParser(network, logger)
with open('resnet50.onnx', 'rb') as f:
parser.parse(f.read())
config = builder.create_builder_config()
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30) # 1GB
engine = builder.build_serialized_network(network, config)
with open('resnet50.plan', 'wb') as f:
f.write(engine)
""")7.3 OpenVINO 转换
python
print()
print("=" * 60)
print("OpenVINO 转换")
print("=" * 60)
print("""
OpenVINO 模型转换：
命令行方式：
# 安装 OpenVINO
pip install openvino-dev
# 转换 ONNX
ovc resnet50.onnx
# 转换 TensorFlow SavedModel
ovc <path_to_saved_model>
# 转换 PyTorch（通过 ONNX）
# 先转为 ONNX，然后：
ovc resnet50.onnx --output_model resnet50.xml
Python API：
from openvino.runtime import Core
# 转换模型
ie = Core()
model = ie.read_model('resnet50.onnx')
# 编译到指定设备
compiled_model = ie.compile_model(model, 'CPU') # 或 'GPU', 'MYRIAD'
# 推理
infer_request = compiled_model.create_infer_request()
output = infer_request.infer({'input': input_data})
""")第8节：模型格式对比表
| 格式 | 框架 | 目标平台 | 优化 | 算子支持 | 动态 Shape |
|---|---|---|---|---|---|
| ONNX | 跨框架 | 通用 | 中等 | 完整 | ✅ 支持 |
| TFLite | TF | 移动端 | 量化优化 | 有限 | ❌ 静态 |
| SavedModel | TF | TF Serving | 一般 | 完整 | ✅ 支持 |
| TensorRT | NVIDIA | GPU | 高度优化 | NVIDIA 优化 | ⚠️ 有限 |
| OpenVINO | Intel | CPU/GPU | 高度优化 | Intel 优化 | ⚠️ 有限 |
| TorchScript | PyTorch | PyTorch | 一般 | 完整 | ⚠️ 有限 |
| FlatBuffer | TFLite | 移动端 | 量化优化 | 有限 | ❌ 静态 |
| IREE | 跨框架 | 通用 | 多级优化 | 可扩展 | ✅ 支持 |
升华
┌──────────────────────────────────────────────────────────────────────────────┐
│ 📚 前端格式：核心原则 │
├──────────────────────────────────────────────────────────────────────────────┤
│ │
│ 1. **格式选择由部署目标决定**：GPU 用 TensorRT，移动端用 TFLite，通用用 ONNX │
│ │
│ 2. **opset 版本是转换质量的关键**：选对版本，避免用模拟算子 │
│ │
│ 3. **动态 Shape 需要显式配置**：不指定 dynamic_axes 就是静态 Shape │
│ │
│ 4. **自定义算子需要处理方案**：注册 symbolic、替换等价算子、或用自定义算子 │
│ │
└──────────────────────────────────────────────────────────────────────────────┘"AI 可查 vs 必须理解"清单
必须理解（不理解就等于不会）：
- 🔴 为什么需要这么多模型格式：每个格式解决什么问题，适用场景是什么
- 🔴 ONNX 的核心概念：Graph、Node、Tensor、opset 版本
- 🔴 PyTorch 导出的三条路径：TorchScript、ONNX、torch.export 的优缺点
- 🔴 动态 Shape 的处理方法：dynamic_axes、多次 Tracing、symbolic_trace
- 🔴 常见转换陷阱：opset 版本、自定义算子、控制流丢失
AI 可查（知道去哪查就行）：
- ✅ 特定格式的 API：torch.onnx.export 参数、TensorRT API 细节
- ✅ 特定硬件的优化建议：TensorRT optimization guide、OpenVINO best practices
- ✅ 特定 opset 版本的算子列表：ONNX OperatorSchemas
- ✅ 特定框架的转换命令：trtexec 参数、ovc 选项
- ✅ 自定义算子的注册方法：symbolic function 的具体实现
学习状态：🟡 开始学习