4. PyTorch模型训练基础：数据集准备（MVTec AD）、分类/分割模型训练、模型保存与导出
好，咱们进入实战环节。前面聊了那么多环境搭建和理论，现在终于要动手训练模型了。这一章，我会带你走一遍完整的流程：从拿到工业缺陷数据集，到训练一个能用的分类或分割模型，最后把模型保存下来准备部署。
说实话，工业缺陷检测和学术界的图像分类不太一样。学术界讲究刷榜，工业界讲究稳定和落地。我刚开始做这个方向时，就踩过不少坑——比如数据集格式不对、训练到一半显存炸了、模型保存后加载报错……嗯，这些今天都会讲到。
4.1 数据集准备：MVTec AD
工业缺陷检测领域，MVTec AD 是绕不开的数据集。它包含了15类工业产品，比如瓶盖、电缆、皮革等。每类都有正常样本和缺陷样本。
为什么选它？ 因为它是目前工业异常检测的标杆数据集。我建议你拿它做练手，效果好了再换自己的数据。
4.1.1 数据集结构
下载下来后，你会发现它的目录结构是这样的：
mvtec_anomaly_detection/
├── bottle/
│ ├── ground_truth/
│ │ ├── broken_large/
│ │ └── broken_small/
│ ├── test/
│ │ ├── broken_large/
│ │ ├── broken_small/
│ │ └── good/
│ └── train/
│ └── good/
├── cable/
│ └── ...
└── ...
每个类别下分三个文件夹：
- train/good：只有正常样本，用于训练
- test/good：正常测试样本
- test/缺陷类别：各类缺陷样本
- ground_truth/缺陷类别：像素级标注（分割任务用）
关键点：训练集只有正常样本。这是工业场景的常见设定——缺陷样本太少，没法用常规分类训练。所以很多工业检测模型都是基于「异常检测」思路做的。
4.1.2 用PyTorch加载数据
我个人习惯用 torchvision.datasets.ImageFolder 来加载。但MVTec AD的目录结构有点特殊，需要自己写个Dataset类。
下面是我常用的代码模板：
import torch
from torch.utils.data import Dataset
from PIL import Image
import os
class MVTecDataset(Dataset):
def __init__(self, root_dir, category, split='train', transform=None):
self.root_dir = root_dir
self.category = category
self.split = split
self.transform = transform
# 构建路径
if split == 'train':
self.img_dir = os.path.join(root_dir, category, 'train', 'good')
else:
self.img_dir = os.path.join(root_dir, category, 'test')
self.img_paths = []
for root, dirs, files in os.walk(self.img_dir):
for f in files:
if f.endswith('.png'):
self.img_paths.append(os.path.join(root, f))
def __len__(self):
return len(self.img_paths)
def __getitem__(self, idx):
img_path = self.img_paths[idx]
image = Image.open(img_path).convert('RGB')
# 从路径中提取标签
label = 0 if 'good' in img_path else 1
if self.transform:
image = self.transform(image)
return image, label, img_path
小技巧：我习惯在返回时带上 img_path，方便调试时定位是哪张图出了问题。这在工业场景下特别有用。
4.2 分类模型训练
分类任务相对简单。我们用ResNet18做骨干网络，把最后一层改成二分类（正常/缺陷）。
4.2.1 模型定义
import torch.nn as nn
import torchvision.models as models
class DefectClassifier(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预训练的特征能帮你快速收敛。我在项目中试过从零训练，效果差很多。
4.2.2 训练循环
训练代码其实很标准，但有几个细节要注意：
def train_one_epoch(model, dataloader, optimizer, criterion, device):
model.train()
total_loss = 0
correct = 0
for images, labels, _ in dataloader:
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)
correct += predicted.eq(labels).sum().item()
return total_loss / len(dataloader), correct / len(dataloader.dataset)
注意：工业缺陷检测中，正常样本和缺陷样本数量往往严重不平衡。我曾经遇到过一个项目，正常样本有5000张，缺陷只有50张。这时候直接用交叉熵损失，模型会学成「永远预测正常」。解决办法是加权重或者用Focal Loss。
4.3 分割模型训练
分割任务更精细，需要输出每个像素的类别。工业上常用U-Net或DeepLabV3。
4.3.1 数据加载（带mask）
分割任务需要同时加载原图和标注图：
class MVTecSegDataset(Dataset):
def __init__(self, root_dir, category, transform=None, mask_transform=None):
self.root_dir = root_dir
self.category = category
self.transform = transform
self.mask_transform = mask_transform
# 只加载测试集（因为只有测试集有标注）
self.img_dir = os.path.join(root_dir, category, 'test')
self.mask_dir = os.path.join(root_dir, category, 'ground_truth')
self.samples = []
for defect_type in os.listdir(self.img_dir):
if defect_type == 'good':
continue
defect_dir = os.path.join(self.img_dir, defect_type)
mask_defect_dir = os.path.join(self.mask_dir, defect_type)
for f in os.listdir(defect_dir):
if f.endswith('.png'):
self.samples.append((
os.path.join(defect_dir, f),
os.path.join(mask_defect_dir, f)
))
def __getitem__(self, idx):
img_path, mask_path = self.samples[idx]
image = Image.open(img_path).convert('RGB')
mask = Image.open(mask_path).convert('L')
if self.transform:
image = self.transform(image)
if self.mask_transform:
mask = self.mask_transform(mask)
return image, mask
4.3.2 训练分割模型
分割模型我用的是轻量化的U-Net，方便后续部署：
import segmentation_models_pytorch as smp
model = smp.Unet(
encoder_name='resnet18',
encoder_weights='imagenet',
in_channels=3,
classes=1, # 二值分割：缺陷/背景
)
criterion = nn.BCEWithLogitsLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
经验之谈：分割任务中，损失函数的选择很关键。BCE Loss对像素级不平衡敏感，我建议配合Dice Loss一起用。公式是：Loss = 0.5 * BCE + 0.5 * Dice。这样既能保证收敛，又能提升分割精度。
4.4 模型保存与导出
训练完了，模型得保存下来。这里我踩过不少坑，给你总结一下。
4.4.1 保存checkpoint
我建议保存完整的checkpoint，包括模型参数、优化器状态、epoch等信息：
checkpoint = {
'epoch': epoch,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'loss': loss,
'accuracy': accuracy
}
torch.save(checkpoint, f'checkpoint_epoch_{epoch}.pth')
为什么这么麻烦？因为训练可能中断。我曾经在训练一个分割模型时，跑了30个小时，结果第28小时断电了。如果只保存了模型参数，恢复训练时优化器状态丢了，学习率调度也得重新来。嗯，从那以后我再也不偷懒了。
4.4.2 导出为TorchScript
部署时，我们通常用TorchScript格式。它不依赖Python环境，C++也能调用：
model.eval()
example_input = torch.randn(1, 3, 224, 224)
traced_model = torch.jit.trace(model, example_input)
traced_model.save('defect_model.pt')
注意：如果你的模型里有动态控制流（比如if语句），用 torch.jit.script 而不是 trace。trace只能记录一次执行路径，遇到不同输入走不同分支时会出错。
4.4.3 导出为ONNX
如果目标平台是NVIDIA Jetson或者用TensorRT加速，ONNX是更好的选择：
import torch.onnx
dummy_input = torch.randn(1, 3, 224, 224)
torch.onnx.export(
model,
dummy_input,
'defect_model.onnx',
input_names=['input'],
output_names=['output'],
dynamic_axes={'input': {0: 'batch_size'}, 'output': {0: 'batch_size'}}
)
这里我设置了动态batch，方便推理时灵活调整。不过要注意，有些算子ONNX不支持，导出前最好用 torch.onnx.export 的 check 参数验证一下。
4.5 完整训练流程示例
最后，给你一个完整的训练脚本框架。你拿过去改改就能用：
def main():
# 1. 数据准备
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 = MVTecDataset('mvtec_anomaly_detection', 'bottle', 'train', transform)
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
# 2. 模型初始化
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = DefectClassifier().to(device)
# 3. 损失和优化器
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
# 4. 训练循环
for epoch in range(50):
loss, acc = train_one_epoch(model, train_loader, optimizer, criterion, device)
print(f'Epoch {epoch}: Loss={loss:.4f}, Acc={acc:.4f}')
# 5. 保存checkpoint
if epoch % 10 == 0:
torch.save({
'epoch': epoch,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'loss': loss,
}, f'checkpoint_epoch_{epoch}.pth')
# 6. 导出部署格式
model.eval()
traced = torch.jit.trace(model, torch.randn(1, 3, 224, 224).to(device))
traced.save('defect_model_final.pt')
if __name__ == '__main__':
main()
好了，这一章的内容就到这里。你跟着跑一遍，应该就能掌握工业缺陷检测模型训练的全流程。下一章我们会讲模型量化，这是部署到边缘设备的关键一步——说实话，那才是真正考验工程能力的地方。