极客笔记Pytorch KeyError: ‘在state_dict中的意外键 “module.encoder.embedding.weight”‘
在本文中，我们将介绍Pytorch中的KeyError异常，特别是当遇到类似于”unexpected key ‘module.encoder.embedding.weight’ in state_dict”的错误消息时。我们将详细解释这个错误的原因，以及如何解决它。
在Pytorch中，KeyError异常通常是由于模型加载或保存时参数不一致造成的。当我们尝试加载预训练模型的参数时，如果模型的结构或参数名称不匹配，就会引发KeyError异常。
让我们来看一个示例，假设我们有一个预训练的模型，它的结构如下：
import torch
import torch.nn as nn
class MyModel(nn.Module):
def __init__(self):
super(MyModel, self).__init__()
self.encoder = nn.Sequential(
nn.Linear(100, 50),
nn.ReLU(),
nn.Linear(50, 10)
)
self.decoder = nn.Sequential(
nn.Linear(10, 50),
nn.ReLU(),
nn.Linear(50, 100)
)
def forward(self, x):
x = self.encoder(x)
x = self.decoder(x)
return x
model = MyModel()
现在我们想加载预训练的权重到这个模型中。如果我们尝试加载预训练的权重，但是模型结构或参数名称与预训练模型不匹配，就会得到类似于”unexpected key ‘module.encoder.embedding.weight’ in state_dict”的错误消息。
让我们创建一个虚拟的预训练权重，然后尝试加载它们：
pretrained_state_dict = {
'encoder.0.weight': torch.randn(100, 50),
'encoder.0.bias': torch.randn(50),
'encoder.2.weight': torch.randn(50, 10),
'encoder.2.bias': torch.randn(10),
'decoder.0.weight': torch.randn(10, 50),
'decoder.0.bias': torch.randn(50),
'decoder.2.weight': torch.randn(50, 100),
'decoder.2.bias': torch.randn(100),
}
model.load_state_dict(pretrained_state_dict)
运行上述代码将会引发KeyError异常，错误消息为”unexpected key ‘module.encoder.embedding.weight’ in state_dict”。这是因为我们的预训练权重字典中没有”module.encoder.embedding.weight”这个键。
阅读更多：Pytorch 教程
解决方法
要解决这个问题，我们可以采取以下几个步骤：
1. 检查模型结构
首先，我们需要检查模型的结构，确保它与预训练模型中的结构一致。在上面的示例中，我们的模型结构是一个具有两个序列模块(encoder和decoder)的简单模型。确保预训练模型的结构与我们的模型结构相匹配。
2. 重命名参数
如果模型结构是一致的，但是参数的名称不匹配，我们可以尝试重命名预训练模型中的参数，使其与我们的模型参数名称一致。在Pytorch中，可以使用state_dict()方法来获取模型的参数字典，然后使用rename()方法来重命名参数。
下面是一个例子，重命名预训练模型中的参数名称：
pretrained_state_dict = {
'encoder.0.weight': torch.randn(100, 50),
'encoder.0.bias': torch.randn(50),
'encoder.2.weight': torch.randn(50, 10),
'encoder.2.bias': torch.randn(10),
'decoder.0.weight': torch.randn(10, 50),
'decoder.0.bias': torch.randn(50),
'decoder.2.weight': torch.randn(50, 100),
'decoder.2.bias': torch.randn(100),
}
# 重命名预训练模型中的参数名称
renamed_state_dict = {}
for key, value in pretrained_state_dict.items():
new_key = key.replace('encoder.0', 'encoder.0') # 在这个示例中，我们假设只有这一层参数不匹配
renamed_state_dict[new_key] = value
# 加载重命名后的权重到模型中
model.load_state_dict(renamed_state_dict)
通过重命名参数名称，我们可以将预训练模型的参数加载到我们的模型中，而不会触发KeyError异常。
3. 忽略不匹配的参数
如果在预训练模型中有一些参数不在我们的模型中使用，我们可以选择忽略这些不匹配的参数。可以通过传递strict=False参数来实现这一点，如下所示：
model.load_state_dict(pretrained_state_dict, strict=False)
使用strict=False参数可以在加载权重时忽略错误的键，但是在其他方面保持严格，确保正确加载匹配的参数。
总结
在本文中，我们介绍了Pytorch中出现KeyError异常的情况，特别是当遇到类似于”unexpected key ‘module.encoder.embedding.weight’ in state_dict”的错误消息时。我们讨论了可能引发异常的原因，以及解决这个问题的方法。重要的是，我们要确保模型结构和参数名称与预训练模型一致，可以通过重命名参数或者忽略不匹配的参数来解决这个问题。