How to load different part of two pretrained models?

I have two pretrained models, namely, model1 and model2.
lets say i have conv1 weights that are in model1 and conv2 that are in model2.
In order to load the weights of conv1 and conv2, can i do the following?

load_name = os.path.join('model1.pth')
checkpoint = torch.load(load_name)
model.load_state_dict(checkpoint['model'], strict=False)

load_name = os.path.join('model2.pth')
checkpoint = torch.load(load_name)
model.load_state_dict(checkpoint['model'], strict=False)

I guess you cannot load into the same model but you can load into two different models.

1 Like

You can do it but calling the child instead of the parent. this is:

model.conv1.load_state_dict(model1)
model.conv2.load_state_dict(model2)

1 Like