Adapting self-supervision tutorial for 2D case - contrastive loss

[ptvu281]

Discussed in #581

^{Originally posted by ptvu281 February 27, 2022}
I'm trying to adapt the self-supervision tutorial for a 2D case. However, I keep getting a tensor shape mismatch error at the contrastive loss step. I've posted a screenshot of the error here: https://imgur.com/a/AXmWKgB
Looking at the source code, I see that the error could be caused by the shape mismatch between negatives_mask and sim_matrix but I'm not sure how to proceed. The things I have changed to adapt for 2D are:

• At the flatten step during training, change end_dim=4 to end_dim=3
• Change img_size and patch_size to a 2-tuple instead of 3-tuple

Any suggestion on how to proceed is appreciated!!

[Nic-Ma]

Hi @yiheng-wang-nv ,

Could you please help confirm this question and provide some comments?

Thanks in advance.

[yiheng-wang-nv]

Hi @ptvu281 ,

the error means:
negatives_mask.shape[1] = 8, and torch.exp(sim_matrix / temperature_tensor).shape[1] = 4,

where negatives_mask.shape[1] equals to self.batch_size * 2,
torch.exp(sim_matrix / temperature_tensor).shape[1] equals to input.shape[1] + target.shape[1], these two values should be the same.

As shown in this line,
could you please try to print the shape of flat_out_v1 and flat_out_v2, and what is the batch size you set?