在 pytorch model 中获取权重和偏差并将它们复制到另一个 model 中的类似层的正确方法是什么？

Question

I am trying to copy weights from a pretrained model layer by layer into another model of exactly similar structure. The original model gives an accuracy of 94% on a binary image classification problem but the target model is unable to predict, results in predicting only one class for entire test set.

For example, I used this piece of code to manually copy weights from the stem of the pretrained model to the stem of the target:

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

where modelmix is the target and modelSep is the pretrained model. Used a similar snippet for all the other layers. The target model is not working even though I can see the weights are similar for all layers. I am Using pytorch 1.1. Thank you

Answer 1

You can create another model that the parameters names are equal for example:

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))

Then you can load the model1 state_dict to model2, and viceversa with the kwargs strict=False .

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

If you want something more custom, you should go as you mention.