解决PyTorch模型加载时的`Missing key(s) in state_dict`和`Unexpected key(s) in`错误

在PyTorch中，当你尝试加载预训练模型时，可能会遇到两种常见的错误：Missing key(s) in state_dict和Unexpected key(s) in. 这些错误通常是由于模型架构与预训练权重不匹配或状态字典中的键丢失/多余引起的。下面我们将详细解释这些错误的原因，并提供解决方案。

错误原因

1. Missing key(s) in state_dict：这个错误意味着预训练权重中的某些键在当前模型架构中找不到。这通常发生在以下几种情况：

   • 你对模型架构进行了修改，但预训练权重是基于原始架构的。
   • 你尝试加载一个不完全匹配的预训练模型（例如，加载一个为更大模型训练的权重到一个较小的模型）。
   • 在训练过程中，你可能删除了某些层或更改了层的名称。

2. Unexpected key(s) in：这个错误表示当前模型架构中存在预训练权重中没有的键。这通常发生在以下情况：

   • 你向模型添加了新的层或组件，但这些新组件在预训练权重中不存在。
   • 你可能更改了某些层的名称或结构，导致预训练权重与新模型架构不匹配。

解决方案

解决这些错误的方法通常取决于你的具体需求。以下是几种可能的解决方案：

1. 确保模型架构匹配：在加载预训练权重之前，请确保你的模型架构与预训练权重完全匹配。如果进行了任何更改，请重新训练模型或从与当前架构匹配的权重中加载。
2. 删除多余的键：如果你确定预训练权重中的某些键是不需要的，你可以在加载状态字典之前从中删除它们。例如：

state_dict = torch.load('path_to_weights.pth')
# 删除不需要的键
del state_dict['unneeded_key']
model.load_state_dict(state_dict)

1. 忽略缺失的键：如果你只是缺少一些权重，并且希望模型能够使用默认初始化来填充它们，你可以在加载状态字典时设置strict=False。这样，缺失的键将被忽略，而模型将使用新初始化的权重来填充它们。例如：

model.load_state_dict(torch.load('path_to_weights.pth'), strict=False)

1. 自定义加载：如果你希望更精细地控制加载过程，你可以编写自定义函数来处理状态字典中的键。例如，你可以根据键的名称选择性地加载权重，或者为缺失的键提供默认值。

def custom_load_state_dict(model, state_dict):
    model_state_dict = model.state_dict()
    for name, param in state_dict.items():
        if name in model_state_dict:
            model_state_dict[name].copy_(param)
        else:
            print(f'Skipping {name} as it is not found in the model')
    model.load_state_dict(model_state_dict)
# 使用自定义函数加载权重
custom_load_state_dict(model, torch.load('path_to_weights.pth'))

无论选择哪种方法，都要确保加载的权重与你的模型架构相匹配，以确保模型的正确性和性能。

总结：
加载预训练权重时遇到Missing key(s) in state_dict和Unexpected key(s) in错误是很常见的。这些错误通常是由于模型架构不匹配或状态字典中的键丢失/多余引起的。通过确保模型架构匹配、删除多余的键、忽略缺失的键或自定义加载过程，你可以解决这些问题。重要的是要始终注意加载的权重与你的模型架构之间的兼容性，以确保模型的正确性和性能。