如何将自己模型的权重转移到同一网络但最后一层的分类数量不同？

Question

我在 Pytorch 中有自己的网络。它首先针对二元分类器（2 类）进行训练。在 10k epochs 之后，我获得了训练后的权重为10000_model.pth。现在，我想将模型用于使用同一网络的 4 类分类器问题。因此，我想将二进制分类器中所有经过训练的权重转移到 4 类问题，而不需要随机初始化的 lass 层。我怎么能做到？这是我的模型

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 20, 5, 1)
        self.conv2 = nn.Conv2d(20, 50, 5, 1)
        self.conv_classify= nn.Conv2d(50, 2, 1, 1, bias=True) # number of class

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.max_pool2d(x, 2, 2)
        x = F.relu(self.conv2(x))
        x = F.max_pool2d(x, 2, 2)
        x = F.relu(self.conv_classify(x))
        return x

这是我做的

model = Net ()
checkpoint_dict = torch.load('10000_model.pth')        
pretrained_dict = checkpoint_dict['state_dict']
model_dict = model.state_dict()
# 1. filter out unnecessary keys
pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
# 2. overwrite entries in the existing state dict
model_dict.update(pretrained_dict)
# 3. load the new state dict
model.load_state_dict(model_dict)

现在，我必须按名称手动删除 pretrained_dict。

pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
pretrained_dict.pop('conv_classify.weight', None)
pretrained_dict.pop('conv_classify.bias', None)

这意味着pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict} 什么都不做。

怎么了？我正在使用 pytorch 1.0。谢谢

Answer 1

两个网络具有相同的层，因此state_dict 中的键相同，所以确实

pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}

什么都不做。两者之间的区别在于权重张量（它们的形状）而不是它们的名称。换句话说，您可以通过[v.shape for v in model.state_dict().values()] 区分两者，但不能通过model.state_dict().keys() 区分。您的“解决方法”方法是正确的。如果您想减少手动操作，我会使用

merged_dict = {}
for key in model_dict.keys():
    if 'conv_classify' in key: # or perhaps a more complex criterion
        merged_dict[key] = model_dict[key]
    else:
        merged_dict[key] = pretrained_dict[key]