📚 课程目录第4章 PyTorch模型训练与导出

4. PyTorch模型训练与导出

好，到了真正动手的环节了。前面我们聊了数据准备和模型选型，现在咱们要把这些理论落地——用 PyTorch 训练一个缺陷分类模型，然后把它保存成 .pt 文件。这个 .pt 文件，说白了就是后续移植到嵌入式设备的“种子”。

我个人习惯把这一步叫做“模型出厂设置”。你想想看，训练阶段的所有参数、权重、甚至优化器状态，都打包在这个文件里。后面做量化、剪枝、转 ONNX，都得从它开始。

4.1 数据集准备：MVTec AD 的加载与预处理

MVTec AD 是工业缺陷检测的标杆数据集。它包含 15 个类别的物体和纹理，每个类别都有正常样本和多种缺陷样本。我刚开始接触这个数据集时，第一反应是“这数据量也太小了吧？”——每个类别也就几百张图。但后来发现，这正是工业场景的常态：样本少、缺陷种类多。

我的经验： 别小看这几百张图。MVTec AD 的图像分辨率高（最高 1024×1024），缺陷区域小，对模型的特征提取能力要求很高。用这个数据集练手，能帮你提前踩很多坑。

加载 MVTec AD 的代码其实不复杂。我习惯用 torchvision 的 ImageFolder 配合自定义 Dataset 类。下面是我常用的加载方式：

import torch
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms, datasets
import os
from PIL import Image

class MVTecADDataset(Dataset):
    """MVTec AD 数据集加载器，支持按类别和缺陷类型筛选"""
    def __init__(self, root_dir, category='bottle', split='train', transform=None):
        self.root_dir = os.path.join(root_dir, category, split)
        self.transform = transform
        self.classes = ['good', 'defect']
        self.samples = []
        
        # 遍历 good 和 defect 文件夹
        for label, class_name in enumerate(self.classes):
            class_dir = os.path.join(self.root_dir, class_name)
            if not os.path.exists(class_dir):
                continue
            for fname in os.listdir(class_dir):
                if fname.endswith(('.png', '.jpg', '.bmp')):
                    self.samples.append((os.path.join(class_dir, fname), label))
    
    def __len__(self):
        return len(self.samples)
    
    def __getitem__(self, idx):
        img_path, label = self.samples[idx]
        image = Image.open(img_path).convert('RGB')
        if self.transform:
            image = self.transform(image)
        return image, label

# 定义预处理流程
train_transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.RandomHorizontalFlip(p=0.5),
    transforms.RandomRotation(10),
    transforms.ColorJitter(brightness=0.1, contrast=0.1),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], 
                         std=[0.229, 0.224, 0.225])
])

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

# 实例化
train_dataset = MVTecADDataset(
    root_dir='./mvtec_anomaly_detection',
    category='bottle',
    split='train',
    transform=train_transform
)

val_dataset = MVTecADDataset(
    root_dir='./mvtec_anomaly_detection',
    category='bottle',
    split='test',
    transform=val_transform
)

train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, num_workers=4)
val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False, num_workers=4)

注意： MVTec AD 的原始数据中，train 文件夹只有 good 样本，test 文件夹包含 good 和 defect。所以训练时你只能看到正常样本，测试时才会遇到缺陷。这其实模拟了工业场景——正常样本容易收集，缺陷样本很难提前获取。

4.2 模型定义：基于 ResNet18 的缺陷分类器

选 ResNet18 作为 backbone，原因很简单：它轻量、成熟、在嵌入式设备上跑得动。我见过不少团队一上来就用 ResNet50 甚至 EfficientNet-b7，结果模型太大，移植到边缘设备时发现推理速度根本跟不上。

这里我做了个小改动——把最后的全连接层换成 2 分类（good vs defect）。嗯，其实也可以做多分类，把每种缺陷类型单独分出来。但工业场景中，很多时候你只需要知道“有没有问题”，具体是什么问题可以交给后续的定位模型去处理。

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

class DefectClassifier(nn.Module):
    """基于 ResNet18 的缺陷二分类模型"""
    def __init__(self, num_classes=2, pretrained=True):
        super(DefectClassifier, self).__init__()
        # 加载预训练权重
        self.backbone = models.resnet18(pretrained=pretrained)
        # 获取 backbone 的特征维度
        in_features = self.backbone.fc.in_features
        # 替换最后的全连接层
        self.backbone.fc = nn.Sequential(
            nn.Dropout(0.5),
            nn.Linear(in_features, 256),
            nn.ReLU(inplace=True),
            nn.Linear(256, num_classes)
        )
    
    def forward(self, x):
        return self.backbone(x)

# 实例化模型
model = DefectClassifier(num_classes=2, pretrained=True)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = model.to(device)

  关键点： 我保留了预训练权重。虽然 MVTec AD 的图像和 ImageNet 差异很大，但预训练模型已经学会了边缘、纹理、形状等基础特征。在工业缺陷检测中，这些低级特征非常有用。我曾经试过从头训练，结果收敛速度慢了一倍，最终精度还低了 3 个点。

4.3 训练流程：损失函数、优化器与训练循环

训练流程其实很标准，但有几个细节我想强调一下。首先是损失函数——二分类用交叉熵，这没问题。但工业缺陷检测中，正负样本极度不平衡（正常样本远多于缺陷样本）。我建议在损失函数里加上类别权重。

# 计算类别权重（根据训练集样本分布）
from sklearn.utils.class_weight import compute_class_weight
import numpy as np

# 获取训练集所有标签
train_labels = [label for _, label in train_dataset.samples]
class_weights = compute_class_weight(
    class_weight='balanced',
    classes=np.unique(train_labels),
    y=train_labels
)
class_weights = torch.tensor(class_weights, dtype=torch.float).to(device)

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss(weight=class_weights)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4, weight_decay=1e-5)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1)

# 训练循环
num_epochs = 30
best_acc = 0.0

for epoch in range(num_epochs):
    model.train()
    running_loss = 0.0
    correct = 0
    total = 0
    
    for images, labels in train_loader:
        images, labels = images.to(device), labels.to(device)
        
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        
        running_loss += loss.item()
        _, predicted = torch.max(outputs, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()
    
    train_acc = 100 * correct / total
    print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {running_loss/len(train_loader):.4f}, Acc: {train_acc:.2f}%')
    
    # 验证
    model.eval()
    val_correct = 0
    val_total = 0
    with torch.no_grad():
        for images, labels in val_loader:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            _, predicted = torch.max(outputs, 1)
            val_total += labels.size(0)
            val_correct += (predicted == labels).sum().item()
    
    val_acc = 100 * val_correct / val_total
    print(f'Validation Acc: {val_acc:.2f}%')
    
    # 保存最佳模型
    if val_acc > best_acc:
        best_acc = val_acc
        torch.save(model.state_dict(), 'best_model.pth')
    
    scheduler.step()

我的习惯： 训练时我通常会同时保存两个 checkpoint：一个是 best_model.pth（验证集精度最高的），另一个是 last_model.pth（最后一个 epoch 的）。为什么？因为有时候验证集精度最高的模型可能过拟合了，最后一个 epoch 的模型反而泛化更好。两个都留着，后面做对比测试时有用。

4.4 模型导出：保存为 .pt 文件

训练完成后，我们需要把模型保存成 .pt 文件。这里有两种方式：只保存权重（state_dict），或者保存完整模型（包括网络结构）。

我个人更推荐保存完整模型。为什么？因为移植到嵌入式设备时，你通常需要加载模型结构 + 权重，如果只保存权重，你还得在目标设备上重新定义一遍网络结构，很容易出错。

# 方式一：保存完整模型（推荐）
torch.save(model, 'defect_classifier_full.pt')

# 方式二：只保存权重
torch.save(model.state_dict(), 'defect_classifier_weights.pth')

# 方式三：保存 checkpoint（包含训练状态）
checkpoint = {
    'epoch': epoch,
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'best_acc': best_acc,
    'class_weights': class_weights
}
torch.save(checkpoint, 'defect_classifier_checkpoint.pth')

避坑指南： 我曾经在导出 .pt 文件时踩过一个坑——用 torch.save(model, ...) 保存完整模型后，在另一台机器上加载时发现报错，原因是那台机器上没有安装 torchvision。因为完整模型保存时会把网络结构也序列化进去，而网络结构中引用了 torchvision 的 ResNet。所以，如果你要跨环境部署，建议用 torch.jit.script 或 torch.jit.trace 导出 TorchScript 格式，它不依赖原始 Python 代码。

4.5 模型验证：加载 .pt 文件并推理

导出之后，一定要验证一下。我见过太多人导出模型后直接扔给嵌入式团队，结果发现推理结果和训练时完全不一样。原因可能是预处理不一致、设备差异、或者模型保存/加载过程中出了问题。

# 加载完整模型
loaded_model = torch.load('defect_classifier_full.pt', map_location='cpu')
loaded_model.eval()

# 单张图片推理
from PIL import Image
import torchvision.transforms.functional as F

def predict_single_image(image_path, model, transform):
    image = Image.open(image_path).convert('RGB')
    input_tensor = transform(image).unsqueeze(0)  # 增加 batch 维度
    with torch.no_grad():
        output = model(input_tensor)
        probabilities = torch.softmax(output, dim=1)
        predicted_class = torch.argmax(probabilities, dim=1).item()
    return predicted_class, probabilities.squeeze().tolist()

# 测试
class_names = ['good', 'defect']
pred_class, probs = predict_single_image('test_defect.png', loaded_model, val_transform)
print(f'Predicted: {class_names[pred_class]}, Probabilities: {probs}')

4.6 知识结构总览

下面这张图把本章的核心流程串起来了。从数据加载到模型导出，每一步都有对应的输出产物。

4.7 常见问题与避坑指南

问题	原因	解决方案
训练 loss 不下降	学习率过大或过小；数据预处理错误	先用 1e-4 的学习率，观察 loss 曲线；检查归一化参数是否正确
验证精度远低于训练精度	过拟合；数据增强不足	增加 Dropout 比例；添加更多数据增强（如 CutMix、MixUp）
加载 .pt 文件时报错	跨环境缺少依赖；PyTorch 版本不兼容	使用 torch.jit.trace 导出 TorchScript；或在目标环境安装相同版本的 torchvision
推理结果与训练时不一致	预处理流程不一致（如 resize 尺寸、归一化参数）	将预处理流程封装成函数，训练和推理共用同一套代码

  核心总结： 训练一个工业缺陷分类模型，技术本身并不复杂。真正的难点在于：数据不平衡的处理、预训练模型的合理使用、以及模型导出后的验证。我见过太多团队在训练阶段花了很多精力调参，结果导出时随便 save 一下就扔给嵌入式团队，最后发现推理结果对不上——这种问题其实完全可以避免，只要你在导出后做一次完整的验证推理。