📚 课程目录第3章模型导出与ONNX转换

3. 模型导出与ONNX转换

说实话，模型导出这一步，看着简单，坑是真不少。我见过太多团队，模型训得漂漂亮亮，一到导出就卡壳。今天咱们就把这块硬骨头啃下来。

为什么需要ONNX？

你想想看，PyTorch模型跑在GPU上，依赖的是Python运行时和CUDA。但生产环境里，我们想要的是轻量、高效、跨平台的推理。ONNX就是那个中间桥梁——它把模型的计算图固化下来，脱离Python也能跑。

我个人习惯把ONNX看作模型的「二进制可执行文件」。一旦导出成功，后续的TensorRT优化、量化、部署，都基于这个中间表示。

核心原则：ONNX导出不是终点，而是TensorRT优化的起点。导出质量直接决定后续优化效果。

从PyTorch/HuggingFace导出模型

导出方式分两种：torch.onnx.export 和 transformers.onnx。我一般用前者，控制力更强。

基础导出流程

import torch
import torch.onnx

# 加载模型
model = YourModel()
model.eval()

# 构造示例输入
dummy_input = torch.randn(1, 3, 224, 224)

# 导出
torch.onnx.export(
    model,
    dummy_input,
    "model.onnx",
    export_params=True,
    opset_version=17,
    do_constant_folding=True,
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={
        'input': {0: 'batch_size'},
        'output': {0: 'batch_size'}
    }
)

这里有个细节：model.eval() 必须调用。我在项目中遇到过有人忘了这步，结果导出的ONNX里还带着Dropout和BatchNorm的训练行为，推理结果完全不对。

HuggingFace模型导出

对于HuggingFace的Transformer模型，我推荐用 transformers.onnx 工具包：

from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers.onnx import export
from pathlib import Path

model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")

# 配置导出参数
from transformers.onnx import FeaturesManager
model_kind, model_onnx_config = FeaturesManager.check_supported_model_or_raise(model)
onnx_config = model_onnx_config(model.config)

# 导出
export(
    preprocessor=tokenizer,
    model=model,
    config=onnx_config,
    opset=17,
    output=Path("llama2.onnx")
)

我的经验：HuggingFace的导出工具会自动处理tokenizer和模型之间的连接。但注意，它默认导出的是encoder-decoder或decoder-only的完整结构。如果你只需要推理部分（比如只做生成），可以手动裁剪。

ONNX opset选择：不是越新越好

ONNX每个opset版本都会新增或修改算子。选哪个版本？我有个简单原则：

opset版本	特点	推荐场景
13-15	稳定，算子覆盖广	通用模型，兼容性优先
16-17	支持更多动态操作	Transformer、LLM
18+	最新算子，但工具链可能不完善	实验性项目

我个人习惯用opset 17。为什么？因为TensorRT对opset 17的支持最成熟。opset 18虽然新，但有些算子TensorRT还没完全覆盖，导出后可能报错。

注意：不要盲目追新。我曾经为了用某个新算子选了opset 18，结果TensorRT编译时直接报「unsupported operator」，最后不得不回退到17。选opset时，先确认你的推理后端支持哪些版本。

动态shape处理：让模型适应不同输入

静态shape的模型只能处理固定尺寸的输入。但实际场景里，batch size、序列长度都是变化的。动态shape就是解决这个问题的。

配置动态轴

在 torch.onnx.export 中，通过 dynamic_axes 参数指定哪些维度是动态的：

dynamic_axes = {
    'input': {
        0: 'batch_size',    # batch维度动态
        1: 'sequence_length'  # 序列长度动态
    },
    'output': {
        0: 'batch_size',
        1: 'sequence_length'
    }
}

torch.onnx.export(
    model,
    dummy_input,
    "dynamic_model.onnx",
    dynamic_axes=dynamic_axes,
    opset_version=17
)

这里有个关键点：dummy_input 的shape决定了导出的基准shape。比如你设 dummy_input = torch.randn(1, 128)，那么导出的ONNX会以batch=1、seq_len=128为基准，但允许运行时传入不同的值。

动态shape的坑

我在项目中遇到过最头疼的问题：模型里有reshape操作，动态shape下reshape的目标shape可能计算错误。比如：

# 错误示例
x = x.reshape(-1, 64)  # 动态shape下，-1可能算出非整数

# 正确做法
batch_size = x.shape[0]
x = x.reshape(batch_size, -1, 64)

避坑指南：导出前，先用 torch.jit.trace 或 torch.jit.script 验证动态shape是否正常工作。我习惯写个测试脚本，传入不同shape的输入，对比PyTorch和ONNX的输出是否一致。

常见导出错误排查

导出报错是家常便饭。下面是我总结的「高频错误清单」：

错误类型	典型信息	解决方案
算子不支持	Unsupported operator: aten::xxx	替换为ONNX支持的算子，或用 `torch.onnx.symbolic` 自定义映射
动态shape冲突	Dimension mismatch in reshape	检查reshape、view等操作，确保动态维度计算正确
控制流问题	Cannot export a model with dynamic control flow	用 `torch.jit.script` 替代 `torch.jit.trace`，或重写为静态控制流
内存不足	CUDA out of memory	减小dummy_input的尺寸，或导出时使用CPU

排查步骤

简化模型：先导出一个最简单的模型（比如去掉所有自定义层），确认基础流程通顺。
逐层排查：用 torch.onnx.export 的 verbose=True 参数，打印每个算子的导出过程。
对比验证：用 onnxruntime 加载导出的ONNX，和PyTorch的输出做数值对比。

import onnxruntime as ort
import numpy as np

# 加载ONNX
sess = ort.InferenceSession("model.onnx")

# 构造输入
input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)

# 推理
outputs = sess.run(None, {'input': input_data})

# 对比PyTorch输出
with torch.no_grad():
    torch_outputs = model(torch.from_numpy(input_data))

# 检查误差
diff = np.abs(outputs[0] - torch_outputs.numpy())
print(f"Max diff: {diff.max()}")

我曾经踩过的坑：有一次导出LLaMA模型，一直报「Unsupported operator: aten::index_put」。查了半天，发现是模型里用了 tensor[index] = value 这种in-place操作。ONNX不支持in-place，改成 tensor = torch.scatter(...) 就好了。所以，导出前最好把模型里的in-place操作都替换掉。

知识体系总览

下面这张图概括了模型导出到ONNX的完整流程和关键决策点：

导出ONNX只是第一步。接下来，这个ONNX文件会进入TensorRT的优化流水线。但如果你在导出阶段就把opset、动态shape、算子兼容性这些问题处理好，后续的优化会顺畅很多。

嗯，今天就先聊到这儿。记住一句话：导出不是终点，是起点。导出质量决定了你能走多远。