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