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

小咪咪 2024-02-19 08:38 37阅读 0赞

### 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`错误。该错误通常由以下几个原因引起：

1.  **模型结构发生变化**：当我们修改了模型的结构（如添加、删除或修改了某些层）后，模型的键名也会发生变化。如果使用旧的`state_dict`加载新的模型，就会出现键名不匹配的情况，从而导致错误。
2.  **多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`，包含一个卷积神经网络和一个全连接层：

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

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

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

现在，我们希望能够加载之前保存的`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`的使用有以下几个常见的场景：

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