PyTorch state_dict 参数名前缀变化后怎么兼容旧检查点

我在把一个单卡模型接入包装层时，第一次遇到的不是张量形状错误，而是一长串 missing_keys 和 unexpected_keys。回头看，旧检查点里的 encoder.weight 与新模型期望的 module.encoder.weight 指向同一组参数，只是模块命名空间多了一层。

处理这类问题的稳妥顺序是：先以当前模型的 state_dict().keys() 为准判断方向，再只对明确的 module. 前缀做映射；用 strict=False 看清缺失与多余键，确认集合完全对齐后再用 strict=True 加载。

官方文档：https://docs.pytorch.org/docs/stable/notes/serialization.html

先确认前缀变化来自模型包装层

PyTorch 的模块状态通常通过 state_dict 序列化，键名由模块树上的属性路径组成。模型被某个包装模块包住后，外层名字可能成为键名的一部分，因此同一组权重会从 backbone.0.weight 变成 module.backbone.0.weight。

这里的重点不是把字符串里出现的 module. 全部替换掉，而是比较来源键和目标模型键。只有当去掉或补上前缀后，键集合能与目标模型对应，才说明这是可以自动迁移的命名变化。

准备一个可重复的加载入口

先把检查点读到 CPU，再处理常见的外层结构。有些训练脚本直接保存 model.state_dict()，有些脚本会把它放在 state_dict 或 model 字段里。迁移函数可以把这个差异收敛到一个入口，但不应把优化器状态误当成模型参数。

from collections import OrderedDict
import torch


def extract_state_dict(checkpoint):
    # 只接受常见的模型权重字段，避免把 optimizer 等训练状态传给模型。
    if isinstance(checkpoint, dict):
        for field in ("state_dict", "model"):
            value = checkpoint.get(field)
            if isinstance(value, dict):
                return value
    # 也兼容直接保存 model.state_dict() 的情况。
    if isinstance(checkpoint, dict):
        return checkpoint
    raise TypeError("检查点中没有可用的 state_dict")


checkpoint = torch.load("old-checkpoint.pt", map_location="cpu")
old_state = extract_state_dict(checkpoint)
old_state = OrderedDict(old_state)

map_location="cpu" 只是让检查点先落到 CPU，避免读取阶段依赖保存时的设备。真正加载前仍要让键名、张量形状和目标模块结构分别通过检查。

按目标模型键名做前缀映射

最可靠的方向判断来自目标模型本身。下面的函数只尝试三种状态：原键已经匹配、去掉统一的 module. 后匹配、补上统一的 module. 后匹配。三种都不匹配时直接报错，避免把错误的映射带进后续训练。

def align_module_prefix(source_state, target_model):
    # 目标键集合是唯一依据，不按某个单独键名猜测映射方向。
    target_keys = set(target_model.state_dict().keys())
    source_keys = set(source_state.keys())
    prefix = "module."

    if source_keys == target_keys:
        return OrderedDict(source_state), "unchanged"

    # 旧权重带 module.、当前模型未包装时，删除这一层前缀。
    if source_keys and all(key.startswith(prefix) for key in source_keys):
        stripped = OrderedDict(
            (key[len(prefix):], value) for key, value in source_state.items()
        )
        if set(stripped.keys()) == target_keys:
            return stripped, "strip-module"

    # 旧权重未带 module.、当前模型有包装层时，补上这一层前缀。
    added = OrderedDict(
        (key if key.startswith(prefix) else prefix + key, value)
        for key, value in source_state.items()
    )
    if set(added.keys()) == target_keys:
        return added, "add-module"

    raise RuntimeError("state_dict 键集合无法通过 module. 前缀对齐")

这个函数故意比较完整集合，而不是只看第一个键。这样可以挡住混合前缀、拼写变化或层结构变化，保留真正需要人工处理的迁移问题。

用 strict=False 获取迁移诊断

load_state_dict 的严格模式要求输入键与当前模型键完全一致。迁移阶段可以先关闭严格检查，读取返回结果里的 missing_keys 和 unexpected_keys，把“没加载到什么”和“多出来什么”分开看。

mapped_state, mode = align_module_prefix(old_state, model)

# 先收集缺失键和多余键；strict=False 只是诊断阶段的开关。
incompatible = model.load_state_dict(mapped_state, strict=False)
print("prefix_mode:", mode)
print("missing_keys:", incompatible.missing_keys)
print("unexpected_keys:", incompatible.unexpected_keys)

# 迁移脚本中不能把非空诊断当成成功，先让调用方看到具体键名。
if incompatible.missing_keys or incompatible.unexpected_keys:
    raise RuntimeError("检查点仍未与当前模型完全对齐")

如果只看到少数新增层出现在 missing_keys，那可能是模型确实增加了参数；如果几乎所有键同时出现在两类列表里，优先回到前缀方向或检查点外层字段，而不是立刻初始化新模型。

确认无缺失后切回严格加载

当映射后的键集合已经与目标模型一致，可以把最终加载改成 strict=True。这样一旦模型定义、检查点或映射逻辑以后发生变化，任务会在加载处显式失败，而不是带着部分随机初始化的权重继续运行。

# 键集合已经完成诊断后，用严格模式作为最终门槛。
model.load_state_dict(mapped_state, strict=True)
model.eval()

无法仅靠 module. 解决时怎么判断

几个容易混淆的问题

直接删除所有键名里的 module. 可以吗？

不建议。只有当目标模型键集合在删除后完全匹配时才应该这样做；嵌套模块自身也可能合法地叫 module，全局替换会破坏真实层级。

strict=False 能不能作为最终方案？

它适合迁移和诊断，不应掩盖未加载的权重。若文章场景要求完整恢复旧模型，应该在确认两类列表都为空后用 strict=True。

为什么键名对上了仍然加载失败？

键名匹配只解决模块路径问题；参数形状、buffer、dtype 和检查点外层结构仍然独立影响加载结果。先看异常的具体字段，再决定是改模型、做张量转换还是补充迁移代码。

总结一下：把旧检查点当作键值映射来处理，先读取当前模型期望的键集合，再只针对统一的 module. 前缀做映射；用 strict=False 看诊断，用 strict=True 做最终确认。这样既能兼容包装层变化，也不会把缺失权重静默带入后续推理。