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.
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