Logging in Multi - GPU and new on_validation_step function

[42elenz]

Hello everyone!
I am overall a bit unsure about my Lightning functions since I am a beginner. I would love to get some feedback about my functions and if I did it right. I was initially working with an old version and then the whole on_validation_step, on_training_step was breaking my logging code.
Also I am wondering why I have to specify:
"ddp_find_unused_parameters_true"

I don't know why there should be unused parameter.

So this is my Trainer:

trainer = Trainer( max_epochs=args.training.epochs, logger=wandb_logger_pt, #log_every_n_steps = args.logging.log_every_n_steps, callbacks=[checkpoint_callback], accelerator="gpu", precision=16, strategy="ddp_find_unused_parameters_true" )

Next this is my Model.
I actually first wanted to use the validation_step_end to append to my self.validation:step_output but the steps were never called. So I had to do it in on_train_epoch and on_valdidation_epoch_end. Is this correct or will this blow up at some point? Probably I just didnt have enough steps so it would be called.

def choose_contrastive_loss_fct(train_criterion_type, hparams):

    if train_criterion_type == 'remove_fn':
        criterion = RemoveFNLoss(hparams.model.temperature, hparams.model.lambda_0)
    elif train_criterion_type == 'binary_remove_fn':
        criterion = BinaryRemoveFNLoss(hparams.model.temperature, hparams.model.lambda_0)
    elif train_criterion_type == 'supcon':
        criterion = SupConLossCLIP(hparams.model.temperature, hparams.model.lambda_0)
    elif train_criterion_type == 'binary_supcon':
        criterion = BinarySupConCLIPLoss(hparams.model.temperature, hparams.model.lambda_0)
    elif train_criterion_type == 'clip':
        criterion = CLIPLoss(hparams.model.temperature, hparams.model.lambda_0)
    elif train_criterion_type == 'ntxent':
        criterion = NTXentLoss(hparams.model.temperature)

    validation_criterion_type = hparams.model.clip_pretrain_validaton_criterion
    if validation_criterion_type == "clip":
        validation_criterion_type = CLIPLoss(hparams.model.temperature, hparams.model.lambda_0, return_both_logits=True)
    
    return criterion, validation_criterion_type

class Contrastive_Training_Model(MultimodalBasis):
    def __init__(self, hparams, fold=''):
        super().__init__(hparams)

        self.save_hyperparameters(hparams)
        self.train_criterion_type = hparams.model.clip_pretrain_train_criterion
        self.label_type_fn = hparams.data.label_type_false_negative_class
        self.train_criterion, self.validation_criterion = choose_contrastive_loss_fct(self.train_criterion_type, hparams)
        self.debugging = hparams.logging.debug_level
        self.correct_val_ids_file = hparams.logging.correct_val_ids_file
        self.fold = fold
        self.training_step_outputs = []
        self.validation_step_outputs = []
        self.counter = 0

        #sanity_check(self.train_criterion_type, self.label_type_fn)
    
    #This is called automatically in trainer class
    def training_step(self, batch, batch_idx):
        """
        Trains contrastive model
        """
        train_mri = batch['cor_mri'] #can be courrpted depends on the settings
        train_questionnaire = batch['cor_questionnaire'] #can be courrpted depends on the settings
        mri_data = batch['mri']
        questionnaire_data = batch['questionnaire']
        id = batch['id']
        fn_label_class = batch['fn_label_class']
        ds_label_class = batch['ds_label_class']

        if self.label_type_fn == "binary":
            fn_label_class = fn_label_class.bool()
        mri_embeddings_projected = self.forward_mri(train_mri) 
        questionnaire_embeddings_projected = self.forward_quest(train_questionnaire)
        loss, logits, labels = self.train_criterion(mri_embeddings_projected, questionnaire_embeddings_projected, fn_label_class)
        self.training_step_outputs.append({'loss':loss, 'logits': logits, 'labels': labels, "ds_label_class": ds_label_class, 'ID': id, 'mri_embeddings': mri_embeddings_projected, 'questionnaire_embeddings': questionnaire_embeddings_projected, 'questionaire_data': questionnaire_data})
        # self.log(f"multimodal.train.loss", loss) implemente later
        #if len(im_views[0])==self.hparams.batch_size:
        #self.calc_and_log_train_embedding_acc(logits=logits, labels=labels, modality='multimodal')

        return {'loss':loss,
                'logits': logits, 
                "ds_label_class": ds_label_class,
                'fn_label_class': fn_label_class, 
                'ID': id,
                'mri_embeddings': mri_embeddings_projected,
                'questionnaire_embeddings': questionnaire_embeddings_projected,
                'questionaire_data': questionnaire_data,}
    
    def training_step_end(self, training_step_output):
        training_step_output = self.trainer.strategy.reduce(training_step_output)
        self.training_step_outputs.append(training_step_output)
        return training_step_output
    

    def validation_step(self, batch, batch_idx):
        """
        Validates contrastive model
        """
        val_mri = batch['cor_mri']
        val_questionnaire = batch['cor_questionnaire']
        mri_data = batch['mri']
        questionnaire_data = batch['questionnaire']
        id = batch['id']
        fn_label_class = batch['fn_label_class']
        ds_label_class = batch['ds_label_class']

        mri_embeddings_projected = self.forward_mri(val_mri)
        questionnaire_embeddings_projected = self.forward_quest(val_questionnaire)
        loss, logits, quest_logits, labels = self.validation_criterion(mri_embeddings_projected, questionnaire_embeddings_projected, fn_label_class)
        self.validation_step_outputs.append({'loss':loss, 'mri_logits': logits, 'quest_logits': quest_logits,'logits': logits, "ds_label_class": ds_label_class, 'fn_label_class': fn_label_class, 'ID': id, 'mri_embeddings': mri_embeddings_projected, 'questionnaire_embeddings': questionnaire_embeddings_projected,'questionaire_data': questionnaire_data})
        return {'loss':loss,
                'mri_logits': logits,
                'quest_logits': quest_logits,
                'logits': logits,
                "ds_label_class": ds_label_class,
                'fn_label_class': fn_label_class, 
                'ID': id,
                'mri_embeddings': mri_embeddings_projected,
                'questionnaire_embeddings': questionnaire_embeddings_projected,
                'questionaire_data': questionnaire_data,}

    def validation_step_end(self, validation_step_output):
        print("validation_step_output", validation_step_output)
        validation_step_output = self.trainer.strategy.reduce(validation_step_output)
        self.validation_step_outputs.append(validation_step_output)

    #At the end of the epcoh. All outputs are in a list.
    def on_train_epoch_end(self):
        train_outputs = self.trainer.strategy.reduce(self.training_step_outputs)
        epoch_loss, epoch_accuracy = evaluation_of_contrastive_outputs(train_outputs,self.debugging, evaluation_type="train")
        fold = self.fold
        self.log(f"cont.train.loss{fold}", epoch_loss, sync_dist=True)
        self.log(f"cont.train.acc{fold}", epoch_accuracy, sync_dist=True)
        self.training_step_outputs.clear()

    def on_validation_epoch_end(self):
        val_outputs = self.trainer.strategy.reduce(self.validation_step_outputs)
        self.counter += 1
        if len(val_outputs) != 0:
            fold = self.fold
            epoch_loss, epoch_accuracy, mri_accuracy, questionnaire_accuracy, mri_quest_accuracy = evaluation_of_contrastive_outputs(val_outputs,self.debugging, evaluation_type="validation", correct_val_ids_file=self.correct_val_ids_file)
            self.log(f"cont.val.loss", epoch_loss, sync_dist=True)
            self.log(f"cont.val.accuracy{fold}", epoch_accuracy, sync_dist=True)
            self.log(f"cont.val.mri_accuracy{fold}", mri_accuracy, sync_dist=True)
            self.log(f"cont.val.questionnaire_accuracy{fold}", questionnaire_accuracy, sync_dist=True)
            self.log(f"cont.val.mri_questionnaire_accuracy{fold}", mri_quest_accuracy, sync_dist=True)
            self.log(f"counter", self.counter, sync_dist=True)
            self.validation_step_outputs.clear()
        
    def configure_optimizers(self):
        optimizer = torch.optim.Adam(
            self.parameters(), 
            lr=self.hparams.training.lr,
            weight_decay=self.hparams.training.weight_decay)
        return optimizer
      

My big problem here is that I want to save the found IDs to a df and save the df. How do I do this with multiple threads? Can I merge threads? How can I just call the root thread for example?

Is it overall correct how I did my logging? I have some parameter that are specific to my task thats the reason I call them on end of the epochs.

For my Downstream-Task I am using a mix of selfmade metrics (balanced Accuracy) and pl metrices.
Does this look ok?:

class ClassificationModel(pl.LightningModule):
    def __init__(
        self,
        hparams,
        num_classes,
        ds_mri_encoder_type,
        model_weights_path=None,
        freezed_backbone=False
    ):
        super().__init__()
        self.save_hyperparameters(hparams)
        input_dim_mri = hparams.model.input_dim_mri
        latent_dim_mri = hparams.model.latent_dim_mri
        mri_encoder_type = hparams.model.mri_encoder_type
        self.debug_level = hparams.logging.debug_level
        self.binary_task = num_classes == 1
        self.correct_val_ids_file = hparams.logging.correct_val_ids_file

        self.mri_encoder = choose_mri_encoder(input_dim_mri, latent_dim_mri, ds_mri_encoder_type, mri_encoder_type, self.debug_level, model_weights_path)
        if freezed_backbone:
            for param in self.mri_encoder.parameters():
                param.requires_grad = False
        self.fc = nn.Linear(latent_dim_mri, num_classes)
        
        if self.binary_task:
            self.loss_fct = nn.BCEWithLogitsLoss()
            self.train_auc = tm.AUROC(task="binary")
            self.val_auc = tm.AUROC(task="binary")
        else:
            self.loss_fct = nn.CrossEntropyLoss()
            self.train_auc = tm.AUROC(num_classes=num_classes, task="multiclass")
            self.val_auc = tm.AUROC(num_classes=num_classes, task="multiclass")

        self.train_accuracy = tm.Accuracy(task="binary" if self.binary_task else "multiclass", num_classes=num_classes)
        self.val_accuracy = tm.Accuracy(task="binary" if self.binary_task else "multiclass", num_classes=num_classes)

        self.train_f1 = tm.F1Score(task="binary" if self.binary_task else "multiclass", num_classes=num_classes)
        self.val_f1 = tm.F1Score(task="binary" if self.binary_task else "multiclass", num_classes=num_classes)

        self.train_balanced_accuracy = 0
        self.val_balanced_accuracy = 0
        self.val_balanced_accuracy_mem = 0

    def forward(self, mri):
        x = self.mri_encoder(mri)
        x = self.fc(x)
        return x
    
    def training_step(self, batch, batch_idx):
        train_mri = batch['mri']
        labels = batch['ds_label_class']

        logits = self.forward(train_mri)
        if self.binary_task:
            labels = labels.unsqueeze(1).float()
        loss = self.loss_fct(logits, labels)

        if self.binary_task:
            logits = torch.sigmoid(logits)
            logits = (logits > 0.5).float()

        self.train_f1.update(logits, labels)
        self.train_accuracy.update(logits, labels)
        self.train_auc.update(logits, labels)

        self.train_balanced_accuracy += balanced_accuracy_score(labels.cpu().numpy(), logits.cpu().numpy())
        self.log('train_loss', loss, on_step=False, on_epoch=True, sync_dist=True)
        return loss
    
    def on_train_epoch_end(self, outputs):
        self.log('train_f1', self.train_f1.compute())
        self.log('train_accuracy', self.train_accuracy.compute())
        self.log('train_auc', self.train_auc.compute())
        final_train_balanced_accuracy = self.train_balanced_accuracy/len(outputs)
        self.log("train_balanced_accuracy", final_train_balanced_accuracy, sync_dist=True)

        self.train_f1.reset()
        self.train_accuracy.reset()
        self.train_auc.reset()
        self.train_balanced_accuracy = 0

    def validation_step(self, batch, batch_idx):
        val_mri = batch['mri']
        labels = batch['ds_label_class']
        ids = batch['id']
        logits = self.forward(val_mri)
        if self.binary_task:
            labels = labels.unsqueeze(1).float()
        loss = self.loss_fct(logits, labels)

        if self.binary_task:
            logits = torch.sigmoid(logits)
            logits = (logits > 0.5).float()
        self.val_f1.update(logits, labels)
        self.val_accuracy.update(logits, labels)
        self.val_auc.update(logits, labels)

        self.val_balanced_accuracy += balanced_accuracy_score(labels.cpu().numpy(), logits.cpu().numpy())
        self.log('val_loss', loss, on_step=False, on_epoch=True, sync_dist=True)
        return {'loss': loss, 'logits': logits, 'labels': labels, 'ID': ids}
    
    def on_validation_epoch_end(self, outputs):
        self.log('val_f1', self.val_f1.compute())
        self.log('val_accuracy', self.val_accuracy.compute())
        self.log('val_auc', self.val_auc.compute())
        final_val_balanced_accuracy = self.val_balanced_accuracy/len(outputs)
        self.log("val_balanced_accuracy", final_val_balanced_accuracy, sync_dist=True)

        self.val_f1.reset()
        self.val_accuracy.reset()
        self.val_auc.reset()
        if self.val_balanced_accuracy_mem < final_val_balanced_accuracy:
            self.val_balanced_accuracy_mem = final_val_balanced_accuracy
            evaluation_of_classifier_outputs(outputs, self.debug_level, binary_task=self.binary_task, correct_val_ids_file=self.correct_val_ids_file)
        self.val_balanced_accuracy = 0

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(
            self.parameters(), 
            lr=self.hparams.training.lr,
            weight_decay=self.hparams.training.weight_decay)
        return optimizer

I hope that you can provide me feedback :)

So the summarization is:

1. Does the overall logging for just metrices I do myself look good? -> will the step function be called? Is it ok to merge all the threads in the epoch end?
2. Why do I have unused parameters?
3. How can I call the root thread or thread 0, or can I wait for all of them?