修改网络后如何有效加载预训练权重

在深度学习的实践中，加载预训练权重是一种常见的做法，它可以帮助我们快速启动项目，缩短训练时间，并可能提升模型的最终性能。然而，当我们在原始网络基础上进行修改后，如何有效地加载这些预训练权重就成了一个需要解决的问题。本文将简明扼要地介绍几种在修改网络后加载预训练权重的方法，并提供实际应用的建议。

一、直接加载预训练权重

如果网络结构的修改较小，且预训练权重中的大部分层仍然存在于新网络中，我们可以尝试直接加载预训练权重。但需要注意的是，如果新网络中存在与预训练权重不匹配的层（如新增层、删除层或层参数变化），直接加载可能会导致错误。

步骤：

1. 使用torch.load()函数加载预训练权重文件。
2. 实例化新网络模型。
3. 尝试使用model.load_state_dict(pretrained_dict, strict=False)加载权重，其中strict=False允许部分不匹配的键存在。

注意：

• strict=False时，不匹配的键将被忽略，但你需要确保这些被忽略的键对应的层对你的模型性能没有显著影响。
• 如果存在大量不匹配，可能需要考虑其他方法。

二、部分加载预训练权重

当网络结构改动较大时，直接加载预训练权重可能不再适用。此时，我们可以选择性地加载预训练权重中与新网络相匹配的层。

步骤：

1. 加载预训练权重文件，并获取其权重字典pretrained_dict。
2. 实例化新网络模型，并获取其权重字典model_dict。
3. 遍历pretrained_dict，检查每个键是否存在于model_dict中，并且对应的权重形状是否一致。
4. 将匹配的权重从pretrained_dict复制到model_dict中。
5. 使用model.load_state_dict(model_dict)加载更新后的权重。

示例代码：

pretrained_dict = torch.load(pretrained_weights_path)
model_dict = model.state_dict()
# 筛选匹配的权重
matched_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict and model_dict[k].shape == v.shape}
# 更新模型权重
model_dict.update(matched_dict)
model.load_state_dict(model_dict)

三、解决常见问题

1. 权重不匹配：如上所述，使用strict=False或手动筛选匹配的权重。
2. 层命名不一致：如果预训练权重和新网络中的层命名不一致，但结构相同，可以在加载权重前对键名进行映射或修改。
3. 新增层初始化：对于新网络中新增的层，需要单独进行初始化。可以使用默认的初始化方法，也可以根据具体任务进行自定义初始化。

四、实际应用建议

1. 评估改动影响：在修改网络结构前，评估改动对模型性能的影响，并确定是否需要加载预训练权重。
2. 备份原始权重：在尝试加载预训练权重前，备份原始权重文件，以防万一加载失败导致数据丢失。
3. 逐步调试：在加载权重过程中，逐步检查每个步骤的输出，确保没有遗漏或错误。

结语

加载预训练权重是深度学习中的一项重要技术，它可以显著提高模型训练效率和性能。然而，在修改网络结构后，如何有效地加载这些权重就成了一个挑战。通过本文介绍的方法，读者可以更加灵活地处理这一问题，并在实际应用中取得更好的效果。