difficulties creating self supervised learning routine

[SantiagoPendon]

Hi to everyone (This is my first post here :) )

I'm trying to pretrain my SwinUNETR from MONAI that uses blocks with shape (32,32,32) with one channel and 32 in batch.
To pretrain it, I want to use SimCLR technique, with the contrastiveLearning Loss function from MONAI. The routine I've made is the following one:

    def forward(self, x_inpaint, X_inpaintBlock, x_1, x_2):

        y_inpaint = self.inpaintingHead(self.decoder(x_inpaint, self.encoder(x_inpaint)))
        y_inpaintBlock = self.inpaintingBlockHead(self.decoder(X_inpaintBlock, self.encoder(X_inpaintBlock)))
        
        z1 = self.contrastiveHead(self.decoder(x_1, self.encoder(x_1))).flatten(start_dim=1,end_dim=4)
        z2 = self.contrastiveHead(self.decoder(x_2, self.encoder(x_2)) ).flatten(start_dim=1,end_dim=4)
        
      
        
    
            
        return y_inpaint, y_inpaintBlock, z1, z2
      
    def training_step(self, batch):
  
       
       
          X, X_inpaint, X_inpaintBlock, mask_inpaint, mask_block  = batch
        
        

           x_1 = utils.applyRandomTransforms(X)
           x_2 = utils.applyRandomTransforms(X)
        
           
           y_inpaint, y_inpaintBlock, z1, z2 = self(X_inpaint, X_inpaintBlock, x_1, x_2) 

            loss, l1in_loss, l1block_loss, loss_contrastive = self.loss(X, y_inpaint,y_inpaintBlock, 
                                                                        mask_inpaint, mask_block,
                                                                         z1, z2)
            
            self.log('loss', loss, prog_bar=True, on_step=True)
            self.log('Inpaint', l1in_loss, prog_bar=True, on_step=True)
            self.log('BlockIn ', l1block_loss, prog_bar=True, on_step=True)
            self.log('Contrastive ', loss_contrastive, prog_bar=True, on_step=True)
            

The loss function fuses the inpainting and the constrastive task:

class comboLoss(torch.nn.Module):
    
    def __init__(self):
        super(comboLoss, self).__init__()
        
        self.maskedL1    = MaskedL1Loss()
        self.ssim  = ssim_loss.SSIMLoss(spatial_dims=3, win_size = 7)
        
        self.constrativeLoss = ContrastiveLoss()
      
    def forward(self, X, y_inpaint, y_inpaintBlock, mask_inpaint, mask_block,z1,z2):

        l1_inpaint     = self.maskedL1(y_inpaint, X, mask_inpaint)
        l1_block       = self.maskedL1(y_inpaintBlock, X, mask_block)
        sim_inpaint = self.ssim(y_inpaint, X) # Predicción y luego Target.
        sim_block   = self.ssim(y_inpaintBlock, X) # Predicción y luego Target.
        
        loss_contrastive = self.constrativeLoss(z1,z2)
        
        loss_inpaint = l1_inpaint + sim_inpaint
        loss_block = l1_block + sim_block
        
        
        total_loss = loss_inpaint + loss_block + loss_contrastive
        
        return total_loss,  loss_inpaint, loss_block, loss_contrastive   

I was looking the augmentations and are OK.

I used the loss with only ssim and l1 and works

In my first attempt I tried using the hidden states from the encoder to create the z1 and z2 but doesn't work well. Finally I decided to follow the MONAI tutorial:
https://github.com/Project-MONAI/tutorials/blob/main/self_supervised_pretraining/vit_unetr_ssl/ssl_train.ipynb

Please, can anybody provide me any suggestion about what could be wrong? (At this moment the model do not perform the contrastive loss)

Thanks a lot.