3. 模型准备：PyTorch模型导出为ONNX、ONNX模型检查与优化、动态轴与静态轴设置

好，咱们进入实战环节的第一个硬骨头——模型准备。

说实话，很多同学在INT8量化上栽跟头，不是量化本身多难，而是模型导出这步就没走稳。我见过太多人拿着一个导坏了的ONNX模型，折腾半天量化精度，最后发现是导出时轴设置错了。嗯，这坑我踩过，所以今天咱们把这步走扎实。

3.1 PyTorch模型导出为ONNX

PyTorch转ONNX，官方给了torch.onnx.export这个接口。用法不复杂，但细节决定成败。

先看一个标准流程：

import torch
import torchvision.models as models

# 加载一个预训练模型
model = models.resnet18(pretrained=True)
model.eval()

# 构造一个 dummy 输入
dummy_input = torch.randn(1, 3, 224, 224)

# 导出
torch.onnx.export(
    model,               # 模型
    dummy_input,         # 示例输入
    "resnet18.onnx",     # 输出路径
    export_params=True,  # 导出参数
    opset_version=11,    # ONNX算子集版本
    do_constant_folding=True,  # 常量折叠优化
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={
        'input': {0: 'batch_size'},
        'output': {0: 'batch_size'}
    }
)

这里有几个关键点，我挨个说。

第一，模型一定要切到eval模式。 我见过有人忘了调model.eval()，结果导出的ONNX里带着BatchNorm的running stats，推理结果直接飘了。这问题排查起来特别隐蔽，我当时花了整整一个下午。

第二，opset_version怎么选？ 我个人习惯用11或13。11版本兼容性最好，几乎所有推理引擎都支持。13版本支持了一些新算子，比如Softmax的改进版。如果你不确定，先用11，遇到不支持的操作再升级。

第三，dynamic_axes是重点。 这个参数决定了你的模型能不能处理变长输入。咱们下一节细讲。

小技巧： 导出时加上verbose=True，可以看到每一步的导出日志。遇到报错时，这个日志能帮你快速定位是哪个算子出了问题。

3.2 ONNX模型检查与优化

导出完了，别急着用。先检查一下模型是不是健康的。

我常用的检查工具有两个：onnx库自带的检查器和onnxruntime的推理验证。

3.2.1 基础检查

import onnx

# 加载模型
model = onnx.load("resnet18.onnx")

# 检查模型结构是否完整
onnx.checker.check_model(model)

print("模型检查通过！")

这个check_model会检查图结构、节点连接、输入输出定义是否合法。如果报错，多半是导出时某些算子没被正确转换。

3.2.2 形状与精度验证

光检查结构不够，还得验证数值对不对。我一般这么做：

import onnxruntime
import numpy as np

# 创建推理会话
ort_session = onnxruntime.InferenceSession("resnet18.onnx")

# 准备输入数据
input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)

# 运行推理
ort_inputs = {ort_session.get_inputs()[0].name: input_data}
ort_outputs = ort_session.run(None, ort_inputs)

# 跟PyTorch结果对比
with torch.no_grad():
    torch_output = model(torch.from_numpy(input_data))

# 计算误差
diff = np.abs(ort_outputs[0] - torch_output.numpy()).max()
print(f"最大误差：{diff:.6f}")

误差在1e-5以内基本没问题。如果误差很大，那就要排查是不是某些算子精度丢失了。

注意： 如果模型里有自定义算子，ONNX可能不认识。这时候需要注册自定义算子，或者用onnx-simplifier做简化。我曾经有个项目用了torch.einsum，导出后ONNX死活不认，最后用onnx-simplifier拆成了基础算子才搞定。

3.2.3 模型优化

ONNX模型可以进一步优化，减少冗余节点。我推荐两个工具：

onnx-simplifier：消除Identity、Cast等冗余节点，合并常量
onnxoptimizer：做图级别的优化，比如节点融合

用法很简单：

# 安装
pip install onnx-simplifier onnxoptimizer

# 简化
python -m onnxsim resnet18.onnx resnet18_sim.onnx

# 优化
python -m onnxoptimizer resnet18_sim.onnx resnet18_opt.onnx

优化后的模型体积可能变小，推理速度也会提升。但注意，优化可能会改变某些算子的行为，优化完最好再跑一遍数值验证。

3.3 动态轴与静态轴设置

这是最容易踩坑的地方，我单独拿出来说。

静态轴：输入输出的形状在编译时固定。比如(1, 3, 224, 224)，batch size永远是1。

动态轴：某些维度可以在运行时变化。比如batch size可以是1、8、32，或者序列长度可以变。

为什么要区分？因为推理引擎对静态形状可以做更多优化，比如内存预分配、算子融合。动态形状则灵活，但性能会打折扣。

看个对比：

特性	静态轴	动态轴
性能	高（可做编译优化）	低（运行时动态调度）
灵活性	低（形状固定）	高（支持变长输入）
内存占用	可预测	可能波动
适用场景	批量推理、固定尺寸	NLP、目标检测

设置动态轴的方法，就是在dynamic_axes里指定哪些维度可变：

dynamic_axes = {
    'input': {
        0: 'batch_size',   # batch维度可变
        2: 'height',       # 高度可变（比如不同分辨率）
        3: 'width'         # 宽度可变
    },
    'output': {
        0: 'batch_size'
    }
}

这里要注意几点：

不是所有维度都能设成动态。 比如卷积层的输入通道数必须是静态的，因为卷积核的权重是固定的。
动态轴越多，推理引擎的压力越大。 我建议只把真正需要变化的维度设成动态，其他保持静态。
有些推理引擎对动态形状支持不好。 比如某些NPU只支持静态形状，这时候你就得在导出时固定所有维度。

  实战建议： 如果你的部署场景里batch size固定（比如服务端一次只处理一个请求），那就用静态轴。如果必须支持变长输入（比如NLP里的不同长度句子），那就只把长度维度设成动态，其他保持静态。别贪多。

3.4 知识体系总览

下面这张图总结了模型准备阶段的核心流程和关键决策点：

这张图把整个流程串起来了。从PyTorch导出开始，经过检查和优化，最后根据你的部署场景决定用静态轴还是动态轴。每一步都有对应的工具和方法。

我个人习惯在导出前先想清楚：我的模型最终要跑在什么硬件上？输入形状会不会变？如果答案是「会变」，那动态轴就是必须的。如果答案是「不会」，那就老老实实用静态轴，省心又高效。

避坑指南： 我曾经在一个语音识别项目里，把时间维度设成了动态轴，结果导出后推理引擎报错说「不支持动态Reshape」。后来查文档才发现，那个引擎只支持batch维度的动态，其他维度必须静态。所以，先看目标推理引擎的文档，再决定动态轴的范围，能省去很多调试时间。

📚 课程目录

3. 模型准备：PyTorch模型导出为ONNX、ONNX模型检查与优化、动态轴与静态轴设置

3.1 PyTorch模型导出为ONNX

3.2 ONNX模型检查与优化

3.2.1 基础检查

3.2.2 形状与精度验证

3.2.3 模型优化

3.3 动态轴与静态轴设置

3.4 知识体系总览