3、PyTorch模型基础：PyTorch模型定义、训练与保存（state_dict、完整模型）、模型评估与导出准备
在工业视觉模型转换链路中，PyTorch是模型训练的起点。无论是缺陷检测、字符识别还是尺寸测量，模型在PyTorch中的定义方式、训练后的保存格式以及评估流程，直接决定了后续ONNX导出和边缘引擎部署的成败。本章聚焦PyTorch模型生命周期中的三个关键环节：模型定义与训练、模型保存（state_dict与完整模型）、模型评估与导出准备，帮助学员建立从训练到部署的工程化思维。
3.1 核心原理：PyTorch模型定义与训练的本质
PyTorch的模型定义遵循面向对象的设计范式。所有神经网络模型必须继承自torch.nn.Module基类，并在__init__方法中声明网络层，在forward方法中定义前向传播逻辑。这种设计使得模型结构、参数管理和计算图构建高度统一。训练过程则基于自动微分机制：每次前向传播构建计算图，反向传播时通过链式法则自动计算梯度，再由优化器更新参数。
在工业视觉场景中，模型定义需要特别关注输入张量的形状。例如，一个用于表面缺陷检测的CNN模型，输入通常是[B, C, H, W]的四维张量，其中B为批次大小，C为通道数（灰度图为1，RGB图为3），H和W为图像高度和宽度。模型定义时需确保卷积层、池化层和全连接层的维度匹配，否则会在训练或推理时报错。
训练循环的核心代码结构如下：
import torch
import torch.nn as nn
import torch.optim as optim
# 模型定义
class DefectCNN(nn.Module):
def __init__(self, num_classes=2):
super(DefectCNN, self).__init__()
self.features = nn.Sequential(
nn.Conv2d(3, 16, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
nn.MaxPool2d(2),
nn.Conv2d(16, 32, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
nn.MaxPool2d(2)
)
self.classifier = nn.Sequential(
nn.AdaptiveAvgPool2d((1, 1)),
nn.Flatten(),
nn.Linear(32, num_classes)
)
def forward(self, x):
x = self.features(x)
x = self.classifier(x)
return x
# 训练循环
model = DefectCNN(num_classes=2)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=1e-3)
for epoch in range(10):
for images, labels in train_loader:
optimizer.zero_grad()
outputs = model(images)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()这段代码展示了模型定义和训练的基本骨架。在工业项目中，还需要加入学习率调度、早停策略、数据增强等机制，但核心逻辑不变。
3.2 判断方法与操作流程：模型保存的两种方式
PyTorch提供了两种模型保存方式，理解它们的区别对后续部署至关重要。
| 保存方式 | API调用 | 保存内容 | 文件大小 | 适用场景 |
|---|---|---|---|---|
| state_dict | torch.save(model.state_dict(), 'model.pth') |
仅模型参数（权重和偏置） | 较小（通常几MB到几十MB） | 推荐用于部署和迁移学习 |
| 完整模型 | torch.save(model, 'model_full.pth') |
模型结构 + 参数 + 优化器状态 | 较大（包含结构定义代码） | 仅用于同环境快速恢复训练 |
操作流程建议：在训练过程中，应定期保存state_dict作为检查点（checkpoint），同时记录当前epoch和最佳验证精度。推荐的文件命名格式为model_epoch{epoch}_acc{accuracy:.4f}.pth，便于回溯。
torch.save({'state_dict': model.state_dict(), 'config': config}, 'checkpoint.pth')的方式打包保存。
3.3 案例拆解：从训练到导出的完整流程
以某PCB板焊点缺陷检测项目为例，项目要求将PyTorch训练的ResNet18模型部署到边缘设备（如Jetson Nano）。初始目标是将模型转换为ONNX格式，并确保推理精度损失小于0.5%。
关键输入数据与参数：
- 训练集：10,000张焊点图像，分辨率224×224，RGB三通道
- 类别：正常、虚焊、桥连、缺件（4类）
- 模型：ResNet18，预训练权重在ImageNet上微调
- 训练超参数：batch_size=32，初始学习率0.001，SGD优化器，动量0.9
操作步骤：
- 模型定义与训练：使用
torchvision.models.resnet18(pretrained=True)加载预训练模型，替换最后一层全连接为4类输出。训练20个epoch，验证精度达到98.2%。 - 保存state_dict：训练结束后，执行
torch.save(model.state_dict(), 'pcb_defect_resnet18.pth')，同时保存类别映射和输入尺寸配置。 - 模型评估：在测试集（2,000张图像）上计算混淆矩阵，确认各类别的精确率和召回率。重点检查虚焊类别的召回率是否满足生产要求（≥95%）。
- 导出准备：将模型切换为评估模式
model.eval()，并包装在torch.no_grad()上下文中，确保BatchNorm和Dropout层行为固定。 - ONNX导出：构造一个虚拟输入
dummy_input = torch.randn(1, 3, 224, 224)，调用torch.onnx.export()导出，设置opset_version=11，并声明动态轴（batch维度）。
结果对比：导出后的ONNX模型在ONNX Runtime上推理，单张图像耗时从PyTorch的12ms降至8ms，精度损失仅0.1%，满足部署要求。
风险与代价：在导出过程中，发现ResNet18中的AdaptiveAvgPool2d在ONNX中映射为动态形状算子，导致部分边缘引擎不支持。解决方案是改用固定尺寸的AvgPool2d，或使用torch.onnx.export的dynamic_axes参数显式声明。
torch.onnx.is_onnx_backend_supported()进行预检。
3.4 常见误区与总结复盘
在PyTorch模型到ONNX的转换过程中，以下误区最为常见：
- 误区一：保存完整模型而非state_dict。完整模型文件包含Python类定义路径，换环境后可能因类不可导入而加载失败。正确做法是保存state_dict，并在加载时重新实例化模型。
- 误区二：忘记切换eval模式。如果模型包含Dropout或BatchNorm层，在训练模式下导出会导致推理结果随机化。必须在导出前调用
model.eval()。 - 误区三：忽略输入尺寸固定。ONNX导出时若未声明动态轴，模型将固定输入尺寸。边缘部署时若图像尺寸变化，需重新导出或使用动态轴。
- 误区四：未验证输出张量。导出后应使用同一输入分别运行PyTorch模型和ONNX模型，对比输出张量的数值差异（通常要求最大误差小于1e-3）。
总结复盘：PyTorch模型基础是工业视觉部署的基石。模型定义阶段要注重结构清晰和参数可配置；训练阶段要关注收敛状态和过拟合风险；保存阶段优先使用state_dict并附带配置信息；评估阶段要全面衡量精度、召回率和推理速度；导出准备阶段则要确保模型处于正确的运行模式，并验证算子兼容性。只有将每个环节做扎实，才能保证后续ONNX转换和边缘引擎部署的顺利进行。