2. 模型训练与导出基础：PyTorch/TensorFlow模型训练要点，导出为ONNX/TorchScript格式

好，咱们进入第二章。这一章，说白了就是讲一件事：怎么把训练好的模型，变成芯片能认的格式。

我见过太多团队，模型在GPU上跑得飞起，一到嵌入式设备就趴窝。为啥？中间那个“导出”环节没做好。你想想看，芯片不认识PyTorch的.pt文件，也不认识TensorFlow的SavedModel。它只认ONNX、TensorRT、或者TFLite这些格式。

所以，这一章咱们就聚焦在训练要点和导出为ONNX/TorchScript这两个核心动作上。我会把我踩过的坑、总结的经验，都揉进去。

2.1 训练阶段的“部署意识”

很多人觉得训练和部署是两拨人干的事。其实不是。我个人习惯在训练脚本里，从一开始就埋下部署的种子。说白了，就是用部署的思维去训练。

2.1.1 算子选择：别用花哨的

PyTorch和TensorFlow都有很多高级算子，比如torch.einsum、tf.map_fn。这些在GPU上跑得挺欢，但导出到ONNX时，经常找不到对应的实现。我遇到过最离谱的一次，一个einsum操作，ONNX直接报“不支持的算子”，折腾了两天才换成矩阵乘法。

所以我的建议是：尽量用基础算子。比如卷积、全连接、ReLU、BatchNorm、残差连接。这些是芯片的“母语”。

2.1.2 输入尺寸：能固定就别动态

芯片推理引擎最怕什么？动态尺寸。你想想看，如果输入图片每次都不一样，芯片就得重新分配内存、重新做算子调度，性能直接打对折。

我建议：训练时就固定输入尺寸。比如224x224、512x512。如果实在需要多尺度，那就训练多个模型，或者用动态轴（后面会讲）。

2.1.3 量化感知训练（QAT）

这个我单独拎出来说。如果你打算在芯片上做INT8量化，那最好在训练阶段就做量化感知训练。说白了，就是让模型在训练时就“习惯”低精度计算。这样导出后再量化，精度损失会小很多。

PyTorch里可以用torch.quantization模块，TensorFlow有tf.quantization。我个人习惯在训练的最后几个epoch开启QAT，效果最好。

2.2 导出为ONNX格式

ONNX是业界最通用的中间格式。几乎所有芯片都支持它。但导出时有很多细节，稍不注意就掉坑里。

2.2.1 PyTorch导出ONNX

PyTorch导出ONNX的核心函数是torch.onnx.export。代码很简单，但参数很重要。

import torch
import torchvision.models as models

model = models.resnet18(pretrained=True)
model.eval()

dummy_input = torch.randn(1, 3, 224, 224)

torch.onnx.export(
    model,
    dummy_input,
    "resnet18.onnx",
    export_params=True,
    opset_version=11,        # 我一般用11或13
    do_constant_folding=True,
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={
        'input': {0: 'batch_size'},   # 动态batch
        'output': {0: 'batch_size'}
    }
)

这里有个关键点：opset_version。不同的opset版本支持的算子不一样。我建议用opset 11或13，兼容性最好。太新的版本，有些老芯片不支持。

2.2.2 TensorFlow导出ONNX

TensorFlow导出ONNX稍微麻烦一点。需要用tf2onnx工具。我个人习惯用命令行方式：

python -m tf2onnx.convert \
    --saved-model ./saved_model \
    --output model.onnx \
    --opset 11 \
    --inputs input:0 \
    --outputs output:0

注意：TensorFlow的模型结构有时候比较复杂，比如有控制流（if/while）。这些在ONNX里可能不支持。我建议导出前先用tf.function把模型包装一下，确保计算图是静态的。

2.2.3 动态轴 vs 静态轴

刚才代码里有个dynamic_axes参数。这个决定了ONNX模型是否支持动态batch。我个人建议：如果芯片支持动态batch，就用动态轴；否则就用静态轴。

静态轴的好处是推理引擎可以做更多编译优化。动态轴的好处是灵活。但有些芯片（比如某些MCU）根本不支持动态轴，你设了也白设。

2.3 导出为TorchScript格式

TorchScript是PyTorch自家的中间格式。它比ONNX更“懂”PyTorch的算子。如果你的目标芯片支持TorchScript（比如某些NVIDIA Jetson设备），那用TorchScript会更省心。

2.3.1 Tracing vs Scripting

TorchScript有两种导出方式：Tracing和Scripting。

• Tracing：给一个示例输入，跟踪模型的计算图。简单、快，但不支持控制流（if/for）。
• Scripting：直接编译模型代码。支持控制流，但需要模型代码是“TorchScript兼容”的。

我个人习惯：能用Tracing就用Tracing。因为Scripting经常遇到代码不兼容的问题，比如用了Python的print或者list.append。

# Tracing 方式
traced_model = torch.jit.trace(model, dummy_input)
traced_model.save("model_traced.pt")

# Scripting 方式
scripted_model = torch.jit.script(model)
scripted_model.save("model_scripted.pt")

2.3.2 TorchScript的算子兼容性

TorchScript虽然比ONNX更“亲民”，但也不是所有PyTorch算子都支持。比如torch.nn.functional.adaptive_avg_pool2d在TorchScript里就有点问题。我建议导出前，先用torch.jit.script试一下，看有没有报错。

2.4 导出后的验证流程

导出不是终点。验证才是。我每次导出后，都会做三件事：

1. 精度验证：用同一批输入，对比原始模型和导出模型的输出。误差应该在1e-5以内。
2. 算子兼容性检查：用onnxruntime或torch.jit跑一遍，看有没有不支持的算子。
3. 性能预估：在芯片上跑一次，看推理时间是否符合预期。如果太慢，可能需要调整模型结构。

验证项	工具	通过标准
精度验证	PyTorch / TensorFlow + ONNX Runtime	输出误差 < 1e-5
算子兼容性	onnxruntime / torch.jit	无报错
性能预估	芯片推理引擎	推理时间 < 目标阈值

嗯，这一章的内容就到这里。记住：训练时多花一分钟考虑部署，导出时就少花一小时填坑。下一章我们会深入芯片推理引擎的细节，看看模型到了芯片上，到底是怎么跑起来的。

---