### PyTorch 中 `load_state_dict` 加载预训练权重后训练速度变慢的原因分析 在 PyTorch 中，使用 `load_state_dict` 方法加载预训练权重是一种常见的实践方式。然而，在某些情况下，加载预训练权重后可能会观察到训练速度显著下降的现象。以下是可能的原因及其对应的解决方案： #### 1. **GPU 和 CPU 的设备一致性** 当模型参数存储在不同的设备（如 GPU 或 CPU）之间时，可能导致性能瓶颈。例如，如果预训练权重是在 GPU 上保存的，而当前模型运行在 CPU 上，则会引发不必要的数据传输开销。 解决方法： 确保模型和权重都位于同一设备上。可以通过显式指定设备来实现这一点： ```python device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model.to(device) pretrained_weights = torch.load("path_to_pretrained.pth", map_location=device) model.load_state_dict(pretrained_weights) ``` #### 2. **梯度计算未被禁用** 即使部分层不需要重新训练，但如果这些层仍然参与反向传播过程，则会导致额外的时间消耗。这通常发生在冻结某些层的情况下。 解决方法： 对于无需更新的层，应将其设置为评估模式并停止梯度跟踪： ```python for param in model.base_model.parameters(): # 假设 base_model 是需要冻结的部分 param.requires_grad = False ``` 此外，还可以调用 `.eval()` 来进一步减少不必要的前向计算中的统计量更新。 #### 3. **BatchNorm 层的状态同步问题** 批量归一化 (Batch Normalization) 层依赖于均值和方差统计数据。加载预训练权重后，如果没有正确处理 BatchNorm 层的行为，可能会导致每次迭代都需要重新计算这些统计信息，从而增加时间成本。 解决方法： 确认 BatchNorm 层是否处于合适的模式下工作。如果是微调阶段，建议保持其默认行为；而在推理阶段则切换至 eval 模式： ```python model.eval() ``` #### 4. **Optimizer 配置不当** 加载预训练权重之后，如果没有适当地初始化优化器或者调整学习率策略，也可能影响整体效率。特别是当新加入的一些随机初始化参数与其他已有的固定参数共存时，这种差异性容易引起收敛困难以及资源浪费。 解决方法： 重新配置 Optimizer 并考虑分组管理不同类型的参数以应用差异化的学习速率： ```python optimizer = torch.optim.Adam([ {'params': model.base_model.parameters(), 'lr': 1e-5}, # 较低的学习率给基础模型 {'params': model.classifier.parameters()} # 默认学习率给分类头 ]) ``` #### 5. **DataLoader 性能不足** 虽然这不是直接由 `load_state_dict` 导致的问题，但在实际项目开发过程中经常忽略 DataLoader 设置的重要性。比如 num_workers 参数过少或 pin_memory 使用不合理都会拖累整个流程的速度表现。 解决方法： 合理配置 Dataloader 参数提升 IO 效率： ```python train_loader = torch.utils.data.DataLoader( dataset=train_dataset, batch_size=64, shuffle=True, num_workers=os.cpu_count(), pin_memory=True ) ``` --- ### 示例代码展示如何高效加载预训练权重并加速后续训练 以下是一个综合以上几点改进措施的例子程序片段： ```python import os import torch from torchvision import models # 定义模型架构 model = models.resnet18(pretrained=False) # 将模型移动到目标设备 device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model.to(device) # 加载预训练权重文件 checkpoint_path = "./resnet18-pretrained.pth" if os.path.exists(checkpoint_path): checkpoint = torch.load(checkpoint_path, map_location=device) model.load_state_dict(checkpoint['state_dict']) # 冻结特征提取部分的参数 for name, param in model.named_parameters(): if "fc" not in name: # 只允许全连接层可训练 param.requires_grad = False # 初始化优化器并对各组件分别设定学习率 optimizer = torch.optim.SGD([ {"params": filter(lambda p: p.requires_grad, model.parameters()), "lr": 0.01}, ]) # 开始正常训练逻辑... ``` --- ###