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?

[Eason-Qin]

Figured out a solution. In this line, the argument parsed into the definition of contrastive loss have 'batch_size=batch_size*2'. However, in the function, the batch size has been multiplied again. The problem was fixed when I set 'batch_size=batch_size' in the definition. If it is the case, I would send a PR to fix that.

[finalelement]

@Eason-Qin I think you might have caught a bug in the SSL tutorial where the batch-size being sent to CL is being multiplied by 2. I am not sure what you are implying by changing the batch size in the definition of the CL. Below is a code snippet, which will work with all batch-sizes.

I would request you to share what code you are trying even if it's a snippet and what batch-size were you setting, so that we can further understand the need for the PR that you are requesting.

@yiheng-wang-nv is right, that both of these need to match up as to what he wrote earlier. The cause for multiplication by 2 within the CL loss function itself is implementation purposes.

import torch
from torch.nn import functional as F
from monai.losses import ContrastiveLoss

def main():

    batch_size = 2
    input = torch.rand(batch_size, 4)
    target = torch.rand(batch_size, 4)

    cl_loss = ContrastiveLoss(batch_size=batch_size, temperature=0.5)
    cl_value = cl_loss(input, target)
    print(cl_value)
    print('Debug here')
    return None

if __name__=="__main__":
    main()