3. 模型训练基础：数据集准备与预处理、模型定义与训练脚本编写、保存训练好的模型（checkpoint）

好，咱们进入实操环节。模型移植这件事，说白了就是让训练好的模型在目标设备上跑起来。但前提是——你得先有一个训练好的模型。这一章，我就带你走一遍模型训练的完整流程：从原始数据到 checkpoint 文件。

我个人习惯把训练流程拆成三个模块：数据流水线、模型与训练逻辑、保存与恢复。这三个模块串起来，就是一个可复用的训练脚本骨架。

3.1 数据集准备与预处理

数据是模型的上限。我在项目中遇到过好几次，模型结构调了半天精度上不去，最后发现是数据预处理跟推理时对不上。嗯，这里要注意：训练时的预处理流水线，必须跟部署时的预处理完全一致。

3.1.1 数据集组织

我建议用文件夹结构来组织数据集，简单直观：

dataset/
├── train/
│   ├── cat/
│   │   ├── cat_001.jpg
│   │   └── cat_002.jpg
│   └── dog/
│       ├── dog_001.jpg
│       └── dog_002.jpg
└── val/
    ├── cat/
    └── dog/

这种结构的好处是：torchvision.datasets.ImageFolder 可以直接读取，省去写数据加载器的麻烦。

3.1.2 预处理流水线

常见的预处理步骤包括：

步骤	操作	说明
1	Resize	统一输入尺寸，常用 224×224 或 256×256
2	RandomCrop / CenterCrop	训练用随机裁剪做数据增强，验证用中心裁剪
3	RandomHorizontalFlip	随机水平翻转，增加数据多样性
4	ToTensor	将 PIL Image 或 numpy 数组转为 Tensor，并缩放到 [0,1]
5	Normalize	用 ImageNet 均值标准差标准化：mean=[0.485,0.456,0.406], std=[0.229,0.224,0.225]

代码实现如下：

from torchvision import transforms

train_transform = transforms.Compose([
    transforms.Resize(256),
    transforms.RandomCrop(224),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225])
])

val_transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225])
])

3.1.3 DataLoader 配置

DataLoader 的参数直接影响训练速度和内存占用：

from torch.utils.data import DataLoader

train_loader = DataLoader(
    train_dataset,
    batch_size=64,
    shuffle=True,
    num_workers=4,      # 数据加载的并行进程数
    pin_memory=True,    # 锁页内存，加速 GPU 传输
    drop_last=True      # 丢弃最后一个不完整的 batch
)

val_loader = DataLoader(
    val_dataset,
    batch_size=64,
    shuffle=False,
    num_workers=4,
    pin_memory=True
)

3.2 模型定义与训练脚本编写

模型定义这块，我习惯用 PyTorch 的 nn.Module 来封装。你想想看，这样后续做量化、导出 ONNX 都方便。

3.2.1 模型定义

以 ResNet-18 为例：

import torch.nn as nn
import torchvision.models as models

class CustomModel(nn.Module):
    def __init__(self, num_classes=2):
        super().__init__()
        # 加载预训练权重
        self.backbone = models.resnet18(pretrained=True)
        # 替换最后一层全连接
        in_features = self.backbone.fc.in_features
        self.backbone.fc = nn.Linear(in_features, num_classes)

    def forward(self, x):
        return self.backbone(x)

这里有个细节：预训练权重。我建议在 ImageNet 上预训练的模型作为起点，尤其是当你的数据集比较小的时候。迁移学习能省很多时间。

3.2.2 训练脚本核心逻辑

训练脚本的骨架其实很固定。我一般这样组织：

import torch
import torch.optim as optim
from torch.optim.lr_scheduler import StepLR

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = CustomModel(num_classes=2).to(device)

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
scheduler = StepLR(optimizer, step_size=30, gamma=0.1)

def train_one_epoch(model, loader, optimizer, criterion, device):
    model.train()
    total_loss = 0.0
    correct = 0
    total = 0

    for images, labels in loader:
        images, labels = images.to(device), labels.to(device)

        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        total_loss += loss.item()
        _, predicted = outputs.max(1)
        total += labels.size(0)
        correct += predicted.eq(labels).sum().item()

    avg_loss = total_loss / len(loader)
    accuracy = 100. * correct / total
    return avg_loss, accuracy

def validate(model, loader, criterion, device):
    model.eval()
    total_loss = 0.0
    correct = 0
    total = 0

    with torch.no_grad():
        for images, labels in loader:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            loss = criterion(outputs, labels)

            total_loss += loss.item()
            _, predicted = outputs.max(1)
            total += labels.size(0)
            correct += predicted.eq(labels).sum().item()

    avg_loss = total_loss / len(loader)
    accuracy = 100. * correct / total
    return avg_loss, accuracy

3.2.3 训练循环

num_epochs = 100
best_acc = 0.0

for epoch in range(1, num_epochs + 1):
    train_loss, train_acc = train_one_epoch(
        model, train_loader, optimizer, criterion, device
    )
    val_loss, val_acc = validate(model, val_loader, criterion, device)
    scheduler.step()

    print(f'Epoch {epoch:3d} | '
          f'Train Loss: {train_loss:.4f} Acc: {train_acc:.2f}% | '
          f'Val Loss: {val_loss:.4f} Acc: {val_acc:.2f}%')

    # 保存最佳模型
    if val_acc > best_acc:
        best_acc = val_acc
        torch.save(model.state_dict(), 'best_model.pth')

为什么会用 state_dict() 而不是整个模型？因为 state_dict() 只保存参数，不保存模型结构。这样更轻量，也方便后续加载到不同设备上。

3.3 保存训练好的模型（Checkpoint）

保存 checkpoint 不只是存个权重文件。我建议保存更多信息，方便训练中断后恢复，或者对比不同 epoch 的效果。

3.3.1 保存完整 Checkpoint

checkpoint = {
    'epoch': epoch,
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'scheduler_state_dict': scheduler.state_dict(),
    'best_acc': best_acc,
    'train_loss': train_loss,
    'val_loss': val_loss
}
torch.save(checkpoint, f'checkpoint_epoch_{epoch}.pth')

3.3.2 加载 Checkpoint 恢复训练

checkpoint = torch.load('checkpoint_epoch_50.pth')
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
start_epoch = checkpoint['epoch'] + 1
best_acc = checkpoint['best_acc']

print(f'Resuming from epoch {start_epoch}, best accuracy: {best_acc:.2f}%')

3.3.3 保存为部署格式

训练完成后，我会额外保存一份 TorchScript 格式，方便后续部署：

model.eval()
example_input = torch.randn(1, 3, 224, 224).to(device)
traced_model = torch.jit.trace(model, example_input)
traced_model.save('model_script.pt')

TorchScript 的好处是：不依赖 Python 环境，可以在 C++ 中直接加载推理。后面做 ONNX 导出也是基于这个格式。

3.4 本章知识结构

我把整个训练流程画成了一张图，方便你对照理解：