Cnn_learner with pretrained head

Is there an easy way to load a pretrained model while retaining the pretrained head, instead of initializing to random weights?

I used custom_head in cnn_learner to solve for now. If you want to only use pytorch, then the approach would be similar.

# Create custom_head
layers = [AdaptiveConcatPool2d(), Flatten(), nn.Linear(4,2)]
head = nn.Sequential(*layers)
    
learn = cnn_learner(data, models.resnet18, custom_head=head)

# Now you can access the head using this
learn.model[1]

# If you want to change it's state dict you can do this
learn.model[1].state_dict()
1 Like

Does doing this include the weights pretrained by models.resnet18? It seems to me this is just creating the architecture?

You are loading a pretrained resnet in the learner.

You can just call models.resnet50(True). That gets you the full Imagenet model with the original head.

2 Likes