How to replace a model's head?

vision
1

I’m trying to implement the simCLR algorithm ([2002.05709] A Simple Framework for Contrastive Learning of Visual Representations) following some steps. One of them is: Create a resnet18 instance (using torchvision) as the base encoder, and replace the classifier head with a projection head. The projection head architecture is as follows: A fully connected layer with 100 units, a relu activation, another fully connected layer with 40 units, a relu activation and finally an output layer with 20 units (so the final features will have dimention=20).

But I got lost with how can I do that in pytorch and would appreciate help. Thanks.

My code:

def create_simclr_model():
  # create new resnet18 as the base encoder
  simclr_model = torchvision.models.resnet18(pretrained=False)
  # change its classifier to a projection head
  # I don't know what to do here.  
  return simclr_model
2

You can replace the last linear layer with an nn.Sequential module containing the described layers:

model = torchvision.models.resnet18(pretrained=False)
model.fc = nn.Sequential(
    nn.Linear(512, 100),
    nn.ReLU(),
    nn.Linear(100, 40),
    nn.ReLU(),
    nn.Linear(40, 20))
out = model(torch.randn(2, 3, 224, 224))
print(out.shape)
> torch.Size([2, 20])
1 Like
3

Hey @algol1
I did something like this to help distinguish between the encoder and the projector.

resnet18 = models.resnet18(
    pretrained=False,
    num_classes=resnet_dim)
encoder = resnet18
projection = nn.Sequential(
    nn.Linear(resnet_dim, hidden_dim),
    nn.ReLU(),
    nn.Linear(hidden_dim, out_dim)
)

This will help you do the evaluation on the SimCLR representation properly.