state_dict 中的意外键：“模型”、“选择”

Question

我目前正在使用 fast.ai 来训练图像分类器模型。

data = ImageDataBunch.single_from_classes(path, classes, ds_tfms=get_transforms(), size=224).normalize(imagenet_stats)
learner = cnn_learner(data, models.resnet34)

learner.model.load_state_dict(
    torch.load('stage-2.pth', map_location="cpu")
)

导致：

torch.load('stage-2.pth', map_location="cpu") 文件 "/usr/local/lib/python3.6/site-packages/torch/nn/modules/module.py", 第 769 行，在 load_state_dict self.class.name, "\n\t".join(error_msgs))) RuntimeError: Error(s) in loading state_dict for Sequential:

...

state_dict 中有意外的键：“model”、“opt”。

我在 SO 中环顾四周并尝试使用以下解决方案：

# original saved file with DataParallel
state_dict = torch.load('stage-2.pth', map_location="cpu")
# create new OrderedDict that does not contain `module.`
from collections import OrderedDict
new_state_dict = OrderedDict()
for k, v in state_dict.items():
    name = k[7:] # remove `module.`
    new_state_dict[name] = v
# load params
learner.model.load_state_dict(new_state_dict)

导致：

RuntimeError: 为 Sequential 加载 state_dict 时出错：

state_dict 中有意外的键：“”。

我正在使用 Google Colab 训练我的模型，然后将训练后的模型移植到 docker 并尝试托管在本地服务器中。

可能是什么问题？会不会是不同版本的pytorch导致模型不匹配？

在我的 docker 配置中：

# Install pytorch and fastai
RUN pip install torch_nightly -f https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html
RUN pip install fastai

虽然我的 Colab 正在使用以下内容：

 !curl -s https://course.fast.ai/setup/colab | bash

Answer 1

我的强烈猜测是stage-2.pth 包含两个顶级项目：模型本身（它的权重）和用于训练它的优化器的最终状态。要仅加载模型，您只需要前者。假设事情是以惯用的 PyTorch 方式完成的，我会尝试

learner.model.load_state_dict(
    torch.load('stage-2.pth', map_location="cpu")['model']
)

更新：在应用我的第一轮建议之后，很明显您正在加载一个保存点创建的模型与您加载它的模型不同（可能配置不同？）。正如您在 pastebin 中看到的那样，保存点包含一些额外层的权重，这些权重在您的模型中不存在，例如 bn3、downsample 等。

“0.4.0.bn3.running_var”、“0.4.0.bn3.num_batches_tracked”、“0.4.0.downsample.0.weight”

同时其他一些键名匹配，但张量的形状不同。

0.5.0.downsample.0.weight 的大小不匹配：从检查点复制形状为 torch.Size([512, 256, 1, 1]) 的参数，当前模型中的形状为 torch.Size([128 , 64, 1, 1])。

我看到一个模式，您一直尝试加载形状为 [2^(x+1), 2^x, 1, 1] 的参数来代替 [2^(x), 2^(x-1), 1, 1]。也许您正在尝试加载不同深度的模型（例如，为 vgg-11 加载 vgg-16 权重？）。无论哪种方式，您都需要弄清楚用于创建保存点的确切架构，然后在加载保存点之前重新创建它。

PS。如果您不确定 - 保存点包含模型权重，以及它们的形状和（自动生成的）名称。它们确实不包含架构本身的完整规范 - 您需要向自己保证，您调用 model.load_state_dict 而 model 与用于创建保存点的架构完全相同.否则你的权重名称可能会不匹配。