模型训练基础回顾：PyTorch/TensorFlow 模型训练流程、模型保存格式与 ONNX 导出
各位同学，欢迎来到第二章。
在正式开始 NPU 推理之前，我们得先把「源头」理清楚。模型是怎么来的？训练完长什么样？怎么才能让 NPU 认识它？
说白了，这一章就是帮你把模型从「训练环境」搬到「部署环境」的准备工作。我见过太多人在这一步踩坑，所以咱们慢慢来，把基础打扎实。
2.1 PyTorch 模型训练流程
PyTorch 的训练流程，我个人觉得是几个框架里最直观的。它基本遵循一个固定的「五步走」套路：
- 定义网络结构：继承
nn.Module，在__init__里搭积木，在forward里定义数据流向。 - 加载数据：用
DataLoader包装数据集，设置 batch size、shuffle 等参数。 - 选择损失函数和优化器：比如
nn.CrossEntropyLoss和optim.Adam。 - 训练循环：遍历每个 epoch，前向传播 → 计算 loss → 反向传播 → 更新权重。
- 保存模型：把训练好的参数存到磁盘上。
嗯，这里要注意一个细节。很多新手会把整个模型对象都存下来，其实没必要。我习惯只保存 state_dict，也就是模型的参数字典。这样更轻量，也更容易跨版本兼容。
# PyTorch 训练核心代码片段
import torch
import torch.nn as nn
import torch.optim as optim
# 1. 定义网络
class MyModel(nn.Module):
def __init__(self):
super().__init__()
self.fc = nn.Linear(784, 10)
def forward(self, x):
return self.fc(x)
model = MyModel()
# 2. 数据加载（略）
# 3. 损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
# 4. 训练循环
for epoch in range(10):
for inputs, labels in dataloader:
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
# 5. 保存模型（推荐方式）
torch.save(model.state_dict(), 'model.pth')model_scripted.pt 的 TorchScript 版本。这样在后续部署到 C++ 环境或某些 NPU 工具链时，会省去很多麻烦。
2.2 TensorFlow 模型训练流程
TensorFlow 这边，我个人觉得它的 API 演进有点快。从 1.x 的 Session 到 2.x 的 Keras，变化挺大。不过现在主流都用 Keras 方式，确实简洁不少。
流程和 PyTorch 大同小异，但写法上有些区别：
- 构建模型：用
tf.keras.Sequential或函数式 API。 - 编译模型：指定优化器、损失函数、评估指标。
- 训练模型：调用
model.fit()。 - 保存模型：可以保存为 H5 格式或 SavedModel 格式。
# TensorFlow Keras 训练核心代码
import tensorflow as tf
# 1. 构建模型
model = tf.keras.Sequential([
tf.keras.layers.Dense(128, activation='relu'),
tf.keras.layers.Dense(10, activation='softmax')
])
# 2. 编译
model.compile(optimizer='adam',
loss='sparse_categorical_crossentropy',
metrics=['accuracy'])
# 3. 训练
model.fit(x_train, y_train, epochs=10, batch_size=32)
# 4. 保存模型
model.save('my_model.h5') # H5 格式
model.save('my_saved_model') # SavedModel 格式（目录）model.save() 默认保存的是 SavedModel 格式（一个文件夹）。如果你指定 .h5 后缀，它才会保存为 H5 文件。但有些 NPU 工具链只认 H5，有些只认 SavedModel。所以导出前，一定要先查清楚目标工具链支持哪种格式。
2.3 模型保存格式对比
你想想看，模型训练完，我们手里可能有 .pth、.h5、SavedModel 这些文件。它们到底有什么区别？
| 格式 | 框架 | 内容 | 特点 |
|---|---|---|---|
.pth |
PyTorch | state_dict（参数字典） | 轻量，只存参数，不存网络结构 |
.h5 |
TensorFlow | 模型结构 + 权重 + 优化器状态 | 单个文件，方便分发 |
SavedModel |
TensorFlow | 模型结构 + 权重 + 签名（Signature） | 目录格式，包含 serving 所需的所有信息 |
我个人建议：如果你只是做实验，.pth 或 .h5 就够了。但如果你要部署到生产环境，尤其是用 TensorFlow Serving，那 SavedModel 是首选。
2.4 模型导出为 ONNX
好了，现在模型有了，但 NPU 不认识 PyTorch 或 TensorFlow 的格式啊。怎么办？
这时候就需要一个「中间人」——ONNX（Open Neural Network Exchange）。它就像一个通用的模型描述语言，把各种框架的模型统一成一种格式。NPU 厂商只要支持 ONNX，就能跑你训练好的模型。
导出 ONNX 其实不复杂，但有几个关键点要注意。
2.4.1 PyTorch 导出 ONNX
import torch
import torch.onnx
# 加载训练好的模型
model = MyModel()
model.load_state_dict(torch.load('model.pth'))
model.eval()
# 创建一个 dummy 输入（形状要和训练时一致）
dummy_input = torch.randn(1, 784)
# 导出 ONNX
torch.onnx.export(
model, # 模型
dummy_input, # 示例输入
'model.onnx', # 输出文件名
input_names=['input'], # 输入节点名称
output_names=['output'], # 输出节点名称
dynamic_axes={'input': {0: 'batch_size'}, # 动态 batch
'output': {0: 'batch_size'}}
)dynamic_axes 参数很重要。它告诉 ONNX 这个模型的 batch size 是可变的。如果不设置，导出的模型就固定了 batch size，比如只能一次推理 1 张图。我在项目中就遇到过这个问题，导出时忘了设动态轴，结果部署时只能一张一张推理，性能大打折扣。
2.4.2 TensorFlow 导出 ONNX
TensorFlow 导出 ONNX 需要借助 tf2onnx 工具。我个人觉得比 PyTorch 稍微麻烦一点，但也不难。
import tf2onnx
import tensorflow as tf
# 加载模型
model = tf.keras.models.load_model('my_model.h5')
# 指定输出路径和输入签名
spec = (tf.TensorSpec((None, 784), tf.float32, name="input"),)
# 转换
model.output_names = ['output']
onnx_model, _ = tf2onnx.convert.from_keras(model, input_signature=spec, opset=13)
# 保存
with open('model.onnx', 'wb') as f:
f.write(onnx_model.SerializeToString())2.5 本章知识结构图
为了让你更直观地理解整个流程，我画了一张图。它展示了从训练到 ONNX 导出的完整链路。
这张图把整个流程串起来了。你可以看到，无论是 PyTorch 还是 TensorFlow，最终都要汇聚到 ONNX 这个中间格式。而 ONNX 模型里包含了计算图、权重、输入输出节点等信息，NPU 工具链就是靠这些信息来解析和优化模型的。