时间戳

首页 > 解决方案 > 在 pytorch 模型中获取权重和偏差并将其复制到另一个模型中的类似层的正确方法是什么？

python - 在 pytorch 模型中获取权重和偏差并将其复制到另一个模型中的类似层的正确方法是什么？

问题描述

我正在尝试将权重从预训练模型逐层复制到另一个结构完全相同的模型中。原始模型在二值图像分类问题上的准确率为 94%，但目标模型无法预测，导致整个测试集只能预测一个类别。

例如，我使用这段代码手动将权重从预训练模型的茎复制到目标的茎：

modelmix.stem[0].weight = modelSep.stem[0].weight
modelmix.stem[1].weight = modelSep.stem[1].weight
modelmix.stem[1].bias = modelSep.stem[1].bias

其中 modelmix 是目标，modelSep 是预训练模型。对所有其他层使用了类似的片段。即使我可以看到所有层的权重都相似，目标模型也不起作用。我正在使用 pytorch 1.1。谢谢

标签： pythonpytorch

解决方案

您可以创建另一个参数名称相同的模型，例如：

import torch.nn as nn

model1 = nn.Sequential()
model1.add_module('layer1', nn.Linear(10, 20))
model1.add_module('layer2', nn.Linear(20, 10))

model2 = nn.Sequential()
model2.add_module('layer1', nn.Linear(10, 20))
model2.add_module('layer2', nn.Linear(20, 10))
model2.add_module('layer3', nn.Linear(10, 5))

然后您可以将 model1 state_dict 加载到 model2，反之亦然 kwargs strict=False。

model2.load_state_dict(model1.state_dict(), strict=False)

如果你想要更定制的东西，你应该按照你说的去。

推荐阅读