SimCLR loss and nearest neighbors computed on hidden features?

[emessori]

Hello,

I was checking your code and noticed that in the definition of SimCLRLoss, more specifically in the forward function, you wrote the following:
input: - features: hidden feature representation of shape [b, 2, dim] output: - loss: loss computed according to SimCLR
From the original SimCLR paper, it was my understanding that the loss function is computed on the output of the projection head (in this case the MLP head attached on top of the backbone resnet). Therefore, I am wondering if I'm missing something or if there is a typo in the comment in SimCLRLoss.

EDIT

There is also something in the computation of the nearest neighbors (function mine_nearest_neighbor, defined for the MemoryBank class in the memory file) which is not clear to me. More specifically, I thought that the nearest neighbors were computed on the 'hidden features', i.e. the vector representation output by the backbone and not by the projection head, but by looking at the dimension of the features it seems that this is not the case. For example, in my config file I set features_dim=128 and since my backbone is ResNet18, I also have backbone['dim'] = 512 as specified both in models\resnet_cifar and models\resnet_stl. However, if I print out the shape of features in mine_nearest_neighbors I can see that each vector is made of 128 entries and not 512.
Does this mean that you compute the nearest neighbors in the space where the loss function is applied and not w.r.t. the outputs of the backbone network?

Thank you and have a nice day!