Unexpected key(s) in state_dict: "module.backbone.bn1.num_batches

Unexpected key(s) in state_dict: "module.backbone.bn1.num_batches_tracked"

在使用PyTorch进行深度学习模型训练和推理时，我们经常会使用state_dict来保存和加载模型的参数。然而，有时当我们尝试加载保存的state_dict时，可能会遇到Unexpected key(s) in state_dict错误，并指明错误的键名。本文将介绍该错误的原因和解决方法。

错误原因

当我们尝试加载模型参数时，state_dict中的键名必须与当前模型中的键名完全匹配。如果不匹配，就会出现Unexpected key(s) in state_dict错误。该错误通常由以下几个原因引起：

模型结构发生变化：当我们修改了模型的结构（如添加、删除或修改了某些层）后，模型的键名也会发生变化。如果使用旧的state_dict加载新的模型，就会出现键名不匹配的情况，从而导致错误。
多GPU训练导致的键名前缀：在使用多GPU进行模型训练时，PyTorch会自动在模型的state_dict中添加前缀module.来表示模型参数来自于不同的GPU。如果我们将单GPU训练的state_dict用于加载多GPU模型，就会出现键名不匹配的情况。

解决方法

以下是几种可能的解决方法：

1. 利用模型的`state_dict`属性名匹配功能

在PyTorch中，可以使用模型的state_dict属性的.keys()方法来查看当前模型的所有键名。然后，我们可以对比保存的state_dict和当前模型的键名，找出不匹配的键名并修改它们。下面是一个示例代码：

pythonCopy code# 加载保存的state_dict
saved_state_dict = torch.load('model.pth')
# 查看当前模型的state_dict键名
model = YourModel()
current_state_dict = model.state_dict()
print("Current model keys:", current_state_dict.keys())
# 修改不匹配的键名
for key in list(saved_state_dict.keys()):
    if key not in current_state_dict:
        new_key = key.replace("module.", "")  # 去除多GPU前缀
        saved_state_dict[new_key] = saved_state_dict.pop(key)
# 加载修改后的state_dict
model.load_state_dict(saved_state_dict)

2. 修改模型代码，适应保存的`state_dict`

如果我们修改了模型的结构，我们可以通过修改模型的代码，使其与保存的state_dict格式相匹配。在加载模型之前，可以先将模型的结构调整为与state_dict结构相同。

3. 使用`torch.nn.DataParallel`进行模型加载

如果模型是使用torch.nn.DataParallel包装的，我们可以使用model = torch.nn.DataParallel(model)来加载模型。这样，模型就可以自动处理多GPU训练导致的键名问题。

pythonCopy codemodel = YourModel()
model = torch.nn.DataParallel(model)  # 加载模型
model.load_state_dict(torch.load('model.pth'))  # 加载state_dict

总结

当加载保存的state_dict时，出现Unexpected key(s) in state_dict错误通常是由于键名不匹配引起的。我们可以通过查看模型的键名和保存的state_dict的键名来找出不匹配的键，并相应地修改它们。另外，使用torch.nn.DataParallel包装模型可以解决多GPU训练导致的键名前缀问题。希望本文能帮助你解决Unexpected key(s) in state_dict错误，并顺利加载模型参数。

示例代码

假设我们有一个图像分类的模型，用于识别猫和狗。我们首先训练了一个模型，并保存了它的state_dict到"model.pth"文件中。然后，我们修改了模型的结构，添加了一个新的全连接层，并希望能够加载之前保存的state_dict。首先，我们定义一个模型类AnimalClassifier，包含一个卷积神经网络和一个全连接层：

pythonCopy codeimport torch
import torch.nn as nn
class AnimalClassifier(nn.Module):
    def __init__(self):
        super(AnimalClassifier, self).__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )
        self.classifier = nn.Sequential(
            nn.Linear(64 * 16 * 16, 256),
            nn.ReLU(inplace=True),
            nn.Linear(256, 2)
        )
    def forward(self, x):
        x = self.features(x)
        x = torch.flatten(x, 1)
        x = self.classifier(x)
        return x

然后，我们训练了模型，并保存了state_dict：

pythonCopy code# 创建模型实例
model = AnimalClassifier()
# 训练模型...
# ...
# 保存state_dict
torch.save(model.state_dict(), 'model.pth')

接下来，我们修改了模型的结构，在全连接层后添加了一个新的ReLU层：

pythonCopy codeimport torch
import torch.nn as nn
class AnimalClassifier(nn.Module):
    def __init__(self):
        super(AnimalClassifier, self).__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )
        self.classifier = nn.Sequential(
            nn.Linear(64 * 16 * 16, 256),
            nn.ReLU(inplace=True),
            nn.Linear(256, 2),
            nn.ReLU(inplace=True)  # 添加新的ReLU层
        )
    def forward(self, x):
        x = self.features(x)
        x = torch.flatten(x, 1)
        x = self.classifier(x)
        return x

现在，我们希望能够加载之前保存的state_dict，并继续训练新的模型。我们可以通过以下代码来加载state_dict并解决键名不匹配的问题：

pythonCopy codeimport torch
import torch.nn as nn
class AnimalClassifier(nn.Module):
    def __init__(self):
        super(AnimalClassifier, self).__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )
        self.classifier = nn.Sequential(
            nn.Linear(64 * 16 * 16, 256),
            nn.ReLU(inplace=True),
            nn.Linear(256, 2),
            nn.ReLU(inplace=True)  # 添加新的ReLU层
        )
    def forward(self, x):
        x = self.features(x)
        x = torch.flatten(x, 1)
        x = self.classifier(x)
        return x
# 创建新的模型实例
model = AnimalClassifier()
# 加载保存的state_dict
saved_state_dict = torch.load('model.pth')
# 查看当前模型的state_dict键名
current_state_dict = model.state_dict()
print("Current model keys:", current_state_dict.keys())
# 修改不匹配的键名
for key in list(saved_state_dict.keys()):
    if key not in current_state_dict:
        new_key = key.replace("classifier.", "classifier.3.") # 修改不匹配的键名
        saved_state_dict[new_key] = saved_state_dict.pop(key)
# 加载修改后的state_dict
model.load_state_dict(saved_state_dict)
# 继续训练新模型...
# ...

通过以上代码，我们成功地加载了之前保存的state_dict，并继续训练了新的模型，同时解决了键名不匹配的问题。

state_dict是PyTorch中用来保存和加载模型参数的一种字典对象。它包含了模型的所有可学习参数的张量（如神经网络的权重和偏置）以及其他相关参数（如优化器的状态），但不包括模型的结构。 state_dict的结构如下：

plaintextCopy code{
    'key1': tensor1,
    'key2': tensor2,
    ...
}

其中，'key' 是一个字符串，对应于模型中的每个参数的名称；'tensor' 是对应于参数的张量。保存模型的state_dict可以通过调用模型的state_dict()方法来获得：

pythonCopy codemodel = MyModel()
...
state_dict = model.state_dict()
torch.save(state_dict, 'model.pth')

加载模型的state_dict可以通过调用torch.load()函数来加载：

pythonCopy codestate_dict = torch.load('model.pth')
model = MyModel()
model.load_state_dict(state_dict)

state_dict的使用有以下几个常见的场景：

保存和加载模型：通过保存和加载state_dict，可以将模型的参数保存到文件并在需要时重新加载参数。
模型的迁移学习和微调：可以将预训练模型的state_dict加载到新模型的对应层中，从而利用预训练模型的参数加快新模型的训练速度或提高性能。
模型参数的共享和复制：可以将一个模型的state_dict复制到另一个模型中，实现参数的共享或复用。
保存和加载优化器状态：优化器的状态信息（如动量、学习率衰减等）通常也存储在模型的state_dict中，可以一同保存和加载。需要注意的是，加载state_dict时，模型的结构应当与保存时的结构完全一致，否则可能会出现加载失败或错误的情况。