4、模型训练与导出：使用PyTorch或TensorFlow训练一个简单图像分类模型（如MNIST），并导出为通用格式（ONNX、TorchScript、SavedModel）

当我们完成了环境准备和数据加载之后，就进入了AI模型落地的第一个关键环节：训练一个可用的模型，并将其导出为可以被推理引擎或生产环境识别的格式。很多初学者在Jupyter Notebook里跑通训练代码后，会误以为“模型已经做好了”，但实际上，Notebook中的模型对象只存在于内存中，一旦进程退出，所有学习到的权重参数就会消失。本章将围绕MNIST手写数字识别这一经典任务，系统讲解使用PyTorch和TensorFlow训练模型的完整流程，并重点剖析三种主流导出格式（ONNX、TorchScript、SavedModel）的适用场景与操作细节。我们会从模型训练的核心原理出发，逐步深入到导出工具链的使用，最后通过一个完整的案例闭环，帮助你建立从“训练脚本”到“可部署产物”的工程化思维。

模型训练与导出工程链路 数据准备 MNIST 28×28灰度 训练集60K / 测试集10K 模型训练 PyTorch / TensorFlow CNN / 全连接网络 训练产物 权重参数 + 计算图 checkpoint / 模型文件 导出格式 ONNX TorchScript SavedModel 部署目标 ONNX Runtime / TensorRT / TFLite / 浏览器 / 移动端 序列化 验证与回滚 精度对比 / 输出一致性检查 / 版本管理 核心要点：训练是过程，导出是桥梁，部署才是目的

上图展示了从数据准备到模型部署的完整工程链路。可以看到，训练得到的模型文件（如PyTorch的.pt或TensorFlow的.h5）只是中间产物，真正要服务于生产环境，必须经过导出（序列化）步骤，将其转换为推理引擎能够高效加载和执行的格式。下面我们分模块深入讲解。

4.1 核心原理：为什么需要导出？训练与推理的差异

在深度学习领域，训练（Training）和推理（Inference）是两个截然不同的阶段。训练阶段需要计算梯度、更新参数，因此必须保留完整的计算图、优化器状态以及中间激活值；而推理阶段只需要前向传播，对延迟和内存占用有极高要求。模型导出，本质上是对训练好的模型进行“瘦身”和“固化”的过程。

具体来说，导出需要解决三个核心问题：格式统一（让不同框架训练的模型能被同一套推理引擎加载）、计算图优化（去除反向传播节点、算子融合、常量折叠）、运行时解耦（不再依赖原始训练框架的Python环境）。以PyTorch为例，其默认的.pt文件保存的是Python对象序列化结果，加载时仍需PyTorch环境；而导出为ONNX后，任何支持ONNX的运行时（如ONNX Runtime、TensorRT）都可以直接加载，无需安装PyTorch。

关键概念区分：

  • Checkpoint（检查点）：保存训练过程中的完整状态，包括模型参数、优化器状态、epoch数等，用于断点续训。
  • Model Artifact（模型产物）：仅包含推理所需的参数和网络结构，用于部署。
  • IR（Intermediate Representation，中间表示）：如ONNX、TorchScript，是框架无关的模型描述格式。

理解这一原理后，我们就能明白为什么导出不是简单的“另存为”，而是需要针对目标平台进行适配和优化的系统工程。

4.2 操作流程：PyTorch与TensorFlow训练及导出实战

下面我们分别演示两个主流框架的完整流程。虽然代码不同，但逻辑主线一致：定义网络 → 加载数据 → 训练循环 → 评估 → 导出。

4.2.1 PyTorch训练MNIST并导出ONNX与TorchScript

PyTorch的导出工具链非常成熟。我们以LeNet-5简化版为例，演示从训练到导出的完整代码。

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader

# 1. 定义网络结构
class SimpleCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.fc1 = nn.Linear(64*5*5, 128)
        self.fc2 = nn.Linear(128, 10)
        self.dropout = nn.Dropout(0.25)

    def forward(self, x):
        x = torch.relu(self.conv1(x))
        x = torch.max_pool2d(x, 2)
        x = torch.relu(self.conv2(x))
        x = torch.max_pool2d(x, 2)
        x = x.view(x.size(0), -1)
        x = torch.relu(self.fc1(x))
        x = self.dropout(x)
        return self.fc2(x)

# 2. 数据加载
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
train_dataset = datasets.MNIST('./data', train=True, download=True, transform=transform)
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)

# 3. 训练配置
model = SimpleCNN()
optimizer = optim.Adam(model.parameters(), lr=0.001)
criterion = nn.CrossEntropyLoss()

# 4. 训练循环（简化版，实际需多个epoch）
model.train()
for batch_idx, (data, target) in enumerate(train_loader):
    optimizer.zero_grad()
    output = model(data)
    loss = criterion(output, target)
    loss.backward()
    optimizer.step()
    if batch_idx % 100 == 0:
        print(f'Batch {batch_idx}, Loss: {loss.item():.4f}')

# 5. 导出为ONNX
model.eval()
dummy_input = torch.randn(1, 1, 28, 28)
torch.onnx.export(model, dummy_input, "mnist_cnn.onnx",
                  input_names=["input"], output_names=["output"],
                  dynamic_axes={"input": {0: "batch_size"}, "output": {0: "batch_size"}})

# 6. 导出为TorchScript（两种方式）
# 方式一：Tracing（跟踪）
scripted_model = torch.jit.trace(model, dummy_input)
scripted_model.save("mnist_cnn_traced.pt")
# 方式二：Scripting（脚本化，支持控制流）
scripted_model_script = torch.jit.script(model)
scripted_model_script.save("mnist_cnn_scripted.pt")

这段代码展示了几个关键点：dynamic_axes参数允许动态batch大小，这在生产环境中非常重要；TorchScript的Tracing方式适合无控制流的网络，而Scripting方式能处理if/for等动态逻辑。

4.2.2 TensorFlow训练MNIST并导出SavedModel

TensorFlow 2.x推荐使用Keras API，导出SavedModel格式最为标准。SavedModel是一个目录结构，包含模型架构、权重和推理签名。

import tensorflow as tf
from tensorflow.keras import layers, models

# 1. 构建模型
model = models.Sequential([
    layers.Conv2D(32, (3, 3), activation='relu', input_shape=(28, 28, 1)),
    layers.MaxPooling2D((2, 2)),
    layers.Conv2D(64, (3, 3), activation='relu'),
    layers.MaxPooling2D((2, 2)),
    layers.Flatten(),
    layers.Dense(128, activation='relu'),
    layers.Dropout(0.25),
    layers.Dense(10, activation='softmax')
])

# 2. 编译与训练
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()
x_train = x_train[..., tf.newaxis] / 255.0
x_test = x_test[..., tf.newaxis] / 255.0
model.fit(x_train, y_train, epochs=3, batch_size=64, validation_split=0.1)

# 3. 导出SavedModel
model.save("mnist_saved_model", save_format="tf")

# 4. 验证导出结果
loaded_model = tf.keras.models.load_model("mnist_saved_model")
test_loss, test_acc = loaded_model.evaluate(x_test, y_test)
print(f"Test accuracy: {test_acc:.4f}")

SavedModel的优势在于其自包含性：目录内的saved_model.pb文件定义了计算图，variables/目录存放权重，assets/存放辅助文件。TensorFlow Serving可以直接加载该目录进行服务部署。

4.3 案例拆解：从训练到ONNX Runtime部署的完整闭环

为了让你更直观地理解整个流程，我们设计一个真实场景：某工业质检项目需要识别MNIST风格的零件编号（10类），要求推理延迟小于5ms，且部署环境只有CPU，没有GPU。

阶段 关键操作 输入/参数 输出/结果
模型训练 使用PyTorch训练SimpleCNN 60K训练样本，batch=128，epoch=5 准确率99.2%，模型大小约1.2MB
格式转换 导出ONNX，开启算子优化 dummy_input=(1,1,28,28) mnist_cnn.onnx，大小1.1MB
推理验证 ONNX Runtime CPU执行 测试集10K，线程数4 平均延迟2.3ms，准确率99.1%
性能调优 开启图优化级别为ORT_ENABLE_ALL session_options.graph_optimization_level 延迟降至1.8ms，吞吐提升22%

在这个案例中，我们遇到了一个典型问题：直接导出的ONNX模型在ONNX Runtime上运行延迟为3.1ms，虽然达标但余量不足。通过分析profiling结果，发现Conv算子占用了60%的时间。我们采取了两个优化措施：一是将batch_size固定为1（因为工业相机是单张触发），避免动态shape带来的额外开销；二是开启ONNX Runtime的图优化级别。最终延迟降至1.8ms，满足产线节拍要求。

风险提示：在导出ONNX时，如果网络中包含自定义算子（如某些注意力机制），需要额外注册算子导出函数。否则导出过程会报错，或者导出的模型在推理时无法识别该算子。建议在项目初期就检查模型是否包含不常见操作。

这个案例的复盘经验是：导出不是终点，验证和调优才是确保生产可用的关键。我们建议在导出后立即进行三件事：精度对比（导出前后输出是否一致）、延迟测试（在目标硬件上跑benchmark）、大小检查（是否适合网络传输或存储）。

4.4 常见误区与总结复盘

在实际教学和项目辅导中，我发现学员在模型导出环节经常犯以下几类错误，这里逐一剖析：

  • 误区一：只保存权重，不保存结构。 例如PyTorch中只使用torch.save(model.state_dict())，加载时需要手动重新定义网络结构。这在部署时非常脆弱，一旦代码重构就会导致模型无法加载。正确做法是保存完整模型或同时保存结构和权重。
  • 误区二：忽略输入输出的预处理/后处理。 训练时的数据归一化参数（均值、标准差）必须与推理时保持一致。很多部署失败案例都是因为推理时忘记做相同的归一化，导致输入分布偏移，准确率骤降。
  • 误区三：盲目追求动态shape。 虽然ONNX支持动态轴，但动态shape会限制推理引擎的优化空间。如果业务场景中batch大小固定，建议在导出时固定shape，以获得更好的性能。
  • 误区四：混淆TorchScript的Tracing和Scripting。 Tracing只能记录实际执行过的路径，如果模型中有依赖数据的控制流（如if x.shape[0] > 10），Tracing会固化该分支，导致推理错误。此时必须使用Scripting。

经验总结：一个稳健的导出流程应该包含以下步骤：① 训练完成后先保存checkpoint；② 加载checkpoint并切换到eval模式；③ 准备dummy input（注意与训练时输入格式完全一致）；④ 执行导出；⑤ 用测试集验证导出模型的精度和输出一致性；⑥ 在目标推理引擎上做性能测试。这六步缺一不可。

此外，关于格式选型，我们给出如下建议：如果团队使用PyTorch且部署环境可控，TorchScript是首选，因为它与PyTorch的兼容性最好；如果需要跨框架或跨硬件（如移动端、嵌入式），ONNX是更通用的选择；如果使用TensorFlow生态，SavedModel配合TensorFlow Serving是最成熟的方案。下表总结了三种格式的核心差异：

对比维度 ONNX TorchScript SavedModel
所属生态 微软+社区（框架中立） PyTorch TensorFlow
运行时支持 ONNX Runtime、TensorRT、TFLite等 LibTorch（C++/Java） TensorFlow Serving、TFLite
动态shape 支持（需配置dynamic_axes） 支持（Scripting模式） 支持（通过SignatureDef）
算子覆盖 受限于ONNX算子集版本 与PyTorch算子对齐 与TF算子对齐
典型场景 跨平台、跨框架部署 PyTorch生态内高性能部署 TF生态、云端服务

最后需要强调的是，模型导出是一个“一次训练、多次适配”的过程。随着业务发展，你可能需要将同一个模型导出为多种格式以适配不同硬件。因此，建议在项目初期就建立模型版本管理机制，记录每次导出的配置参数（如算子集版本、优化级别、输入输出签名），以便追溯和复现。通过本章的学习，你应该已经掌握了从训练到导出的完整方法论，并能够根据实际场景选择合适的导出策略。

← 上一章 📖 返回目录 下一章 →