[FIXED] What's the best way to load two pretrained model partially from state-dict in Pytorch?

Issue

I am trying to load two separately trained models except for the last layer and want to train the last layer separately combining these two models. I defined e new nn.Module class and load these pretrained model inside that class and in the forward path tried to return the value before the last layer.

class New_net(nn.Module):
    def __init__(self):
        super(New_net, self).__init__()
        self.net1 = net1()
        self.net2 = net2()
        self.fc= nn.Linear(512, 2)
        self._initialize_weights()

    def _initialize_weights(self):
        checkpoint = torch.load('save_model/checkpoint_net1.t7')
        self.net1.load_state_dict(checkpoint['state_dict'])

        checkpoint = torch.load('save_model/checkpoint_net2.t7')
        self.net2.load_state_dict(checkpoint['state_dict'])       

    def forward(self, x):
        x1 = self.net1(x)
        x2 = self.net2(x)
        x=torch.cat((x1,x2),dim=1)
        x=self.fc(x)
        return x

but it seems it is not loading the model accurately. What's the correct way to do that

Solution

I figured that. Instead of weight initialization, I did the following

#load net1 model partially
checkpoint = torch.load('save_model/checkpoint_net1.t7')
pretrained_dict=checkpoint['state_dict']

net1_dict=net.net1.state_dict()
pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in net1_dict}
net1_dict.update(pretrained_dict)
net.net1.load_state_dict(net1_dict)

#load net2 model partially
checkpoint = torch.load('save_model/checkpoint_net2.t7')
pretrained_dict=checkpoint['state_dict']
net2_dict=net.net2.state_dict()
pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in net2_dict}
net2_dict.update(pretrained_dict)
net.net2.load_state_dict(net2_dict)

Answered By - Joy Mazumder