Copy weights from pretrained ResNet-50 model to pruned model with different channels

Hey guys, I want to ask if there’s a method to copy weights from pretrained ResNet-50 model to pruned model.

Now I have the prune cfg, which is the number of remain channels in each layer

e.g. pruned_cfg = [18, 23, 24, 25, 20, 30, 46, 46, 30, 39, 49, 71, 52, 59, 130,
109, 62, 84, 73, 85, 86, 92, 112, 110, 105, 106, 24, 17, 41, 37, 37, 46]

And I have masks, which stores which channel are preserved in original models(“1”) and which is being pruned(“0”)

e.g. masks[0]=tensor([1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,
1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 0., 1., 1.,
0., 0., 1., 1., 1., 1., 1., 1., 0., 1., 0., 0., 0., 1., 1., 0., 0., 0.,
1., 1., 1., 1., 1., 0., 1., 1., 1., 0.])

    torch.sum(masks[0])=18 (remaining channels)

Now I want to create a new ResNet-50 model with pruned_cfg, and copy the corresponding channel weights from the original model (copy all 18 corresponding channel weights from the original model in layer1, …)

Is there an easy approach?
Thanks very much!

I don’t know if there is an easy way to perform your desired loading and you would most likely need to load each weight separately:

model = models.resnet18()

mask = torch.tensor([1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,
1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 0., 1., 1.,
0., 0., 1., 1., 1., 1., 1., 1., 0., 1., 0., 0., 0., 1., 1., 0., 0., 0.,
1., 1., 1., 1., 1., 0., 1., 1., 1., 0.]).bool()

sd = models.resnet18(pretrained=True).state_dict()

with torch.no_grad():
    model.conv1 = nn.Conv2d(3, int(mask.sum().item()), 7, 2, 3, bias=False)
    model.conv1.weight.copy_(sd["conv1.weight"][mask])

Depending on the config and how you’ve created it, you might be able to iterate the state_dict keys etc., which could make the loading more convenient.

Sorry, I didn’t see your reply until now, thanks for you answer