第3章模型训练基础

3、模型训练基础：PyTorch/TensorFlow框架简介、训练一个简单的分类模型、模型保存格式（.pt/.pth/.h5）

好，咱们进入第三章。这一章，说白了就是动手干活了。

前面两章我们聊了硬件和模型的关系，但模型到底是怎么来的？总不能凭空变出来吧。对，得训练。而训练，就得选个趁手的框架。

我个人习惯把框架比作「厨具」。PyTorch 像是一把好用的中式炒锅，灵活、直观，适合研究新菜式；TensorFlow 则更像是一台商用烤箱，稳定、标准化，适合大规模出餐。没有绝对的好坏，看你要做什么。

3.1 PyTorch 与 TensorFlow：我该怎么选？

先说说 PyTorch。我记得 2018 年刚接触它时，最大的感受就是「动态图」带来的爽快感。你写一行代码，它立刻就能执行，调试起来跟写普通 Python 脚本一样。这对于搞研究、快速验证想法来说，简直是神器。

举个例子，你想在训练过程中临时打印某个中间层的输出，PyTorch 里直接加个 print 就行。这在 TensorFlow 1.x 时代，你得先建好计算图，再用 session.run() 去拉数据，非常绕。

核心区别一句话：

PyTorch：动态图，调试友好，学术界主流，部署时需转 ONNX 或 TorchScript。
TensorFlow：静态图（2.x 也支持动态了），生产部署成熟，Keras 接口友好，移动端支持好。

不过现在 TensorFlow 2.x 也引入了 Eager Execution，动态图体验已经追平了不少。但我在项目中遇到过一个问题：当你需要做自定义算子或者复杂的控制流时，PyTorch 的 Pythonic 风格依然更顺手。而 TensorFlow 的 SavedModel 格式在服务端部署时，确实比 PyTorch 的 TorchScript 要稳定一些。

我的建议：

如果你刚入门，或者主要做 CV/NLP 研究，从 PyTorch 开始。如果你要落地到移动端或生产环境，且团队有 TF 经验，选 TensorFlow。别纠结，两个都学也不难。

3.2 训练一个简单的分类模型（以 PyTorch 为例）

光说不练假把式。咱们用 PyTorch 训练一个最简单的分类器——识别手写数字 MNIST。这个数据集就像编程界的 "Hello World"，但背后流程跟训练大模型是一样的。

整个流程分四步：

加载数据：把图片和标签准备好。
定义网络：搭一个简单的卷积神经网络。
训练循环：前向传播 -> 计算损失 -> 反向传播 -> 更新权重。
保存模型：把训练好的参数存下来。

来，看代码：

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms

# 1. 加载数据
transform = transforms.Compose([transforms.ToTensor(),
                                transforms.Normalize((0.5,), (0.5,))])
train_set = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
train_loader = torch.utils.data.DataLoader(train_set, batch_size=64, shuffle=True)

# 2. 定义网络（简单 CNN）
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)

    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.fc2(x)
        return x

model = SimpleCNN()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 3. 训练循环
for epoch in range(5):
    for images, labels in train_loader:
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
    print(f'Epoch {epoch+1}, Loss: {loss.item():.4f}')

# 4. 保存模型
torch.save(model.state_dict(), 'mnist_cnn.pth')
print("模型已保存为 mnist_cnn.pth")

嗯，这里要注意：model.state_dict() 保存的是模型的「参数」字典，不包含网络结构。你加载时，需要先实例化一个一模一样的网络，再 load 参数。

我曾经踩过的坑：

有一次我保存了完整的模型（用 torch.save(model, 'model.pt')），结果换了一台机器，Python 版本和 PyTorch 版本不一样，加载时直接报错。从那以后，我坚持只保存 state_dict，再配合一个模型定义脚本。这样兼容性最好。

3.3 模型保存格式：.pt / .pth / .h5 到底有啥区别？

很多新手会被这些后缀搞晕。其实它们只是约定俗成的命名习惯，技术上没有强制规定。但作为工程师，我们得知道背后的逻辑。

后缀	框架	内容	典型用法
.pth	PyTorch	通常只保存 state_dict（参数）	`torch.save(model.state_dict(), 'model.pth')`
.pt	PyTorch	可以保存完整模型或 TorchScript	`torch.save(model, 'model.pt')` 或 `torch.jit.save()`
.h5	TensorFlow/Keras	保存模型架构 + 权重 + 优化器状态	`model.save('model.h5')`

你想想看，为什么会有这么多格式？说白了，就是不同阶段需要不同的东西。

训练阶段：你需要保存 checkpoint，包含参数、优化器状态、epoch 数，方便断点续训。这时候我习惯用 .pth 或 .ckpt。
推理阶段：你只需要参数和网络结构，不需要优化器。PyTorch 推荐用 TorchScript（.pt），TensorFlow 推荐用 SavedModel（文件夹格式）。
跨框架转换：比如 PyTorch 转 TensorFlow，通常先转成 ONNX（.onnx），再转成目标格式。

实战建议：

训练时：保存 .pth 文件，命名带上 epoch 和 loss，比如 model_epoch5_loss0.23.pth。
部署时：导出为 .pt（TorchScript）或 .onnx，确保不依赖 Python 环境。
TensorFlow 用户：直接用 .h5 或 SavedModel，后者更推荐。

我记得有一次部署一个 PyTorch 模型到手机端，对方只接受 TensorFlow Lite 格式。我先把 .pth 转成 ONNX，再转成 TF，最后转成 TFLite。中间踩了不少坑，比如算子不支持、维度不对齐。所以，从一开始就考虑好部署目标，能省很多事。

小技巧：

如果你不确定该用哪种格式，记住一个原则：训练用 .pth，部署用 .onnx 或 TorchScript。.h5 在 Keras 生态里很通用，但跨框架时还是 ONNX 最保险。