Using data from training_epoch_end in validation_epoch_end.

[amdson]

I'm running a metric learning model, and I'd like to use embeddings from the training step for KNN lookups in the validation step. I tried to do so with the following code.

def training_epoch_end(self, outputs):
        embeddings = []
        for output in outputs:
            embeddings.append(output["embeddings"])
        self.training_embeddings = torch.cat(embeddings, dim=0)
        print("Finished training epoch")
        return super().training_epoch_end(outputs)

def validation_epoch_end(self, outputs):
      train_embeddings = self.train_embeddings
      #Use train embeddings...
      self.train_embeddings = None #Delete embeddings for future use.

However, self.train_embeddings is either None, or uninitialized, when validation_epoch_end executes. Is there a way to communicate between these two methods, or send information from both to a final callback?

[amdson]

Just saw other answers explaining the order of train_epoch_end and val_epoch_end hooks. This issue was solved by moving any code needing both validation and training data to "on_train_epoch_end".

[samils7]

Do you perform KNN lookup also in "on_train_epoch_end"?

[amdson]

Yeah, I didn't include the code here, but I stored the model embeddings during training, and did KNN lookups on them in on_train_epoch_end for validation.

[samils7]

on_train_epoch_end function called 4 times when training ddp with 4 gpus. Each calls have 1/4 th of train and val embeddings seperately. Do you have any idea how to collect all embeddings together in order to perform knn lookup?

[akihironitta]

You can gather things (e.g. tensors) from different devices by calling self.all_gather(...) inside your LightningModule.

https://pytorch-lightning.readthedocs.io/en/1.7.7/common/lightning_module.html?highlight=all_gather#all-gather

[samils7]

Thanks Akihiro. I've just tried that simply like:

    def training_epoch_end(self, outputs):
        self.train_embeddings = torch.cat([x['embeddings'] for x in outputs])
        self.train_y = torch.cat([x['y'] for x in outputs])
        print('training_epoch_end! --- train embeddings: ', self.train_embeddings.shape)

    def validation_epoch_end(self, outputs):
        self.val_embeddings = torch.cat([x['embeddings'] for x in outputs])
        self.val_y = torch.cat([x['y'] for x in outputs])
        print('validation_epoch_end! --- val embeddings: ', self.val_embeddings.shape)

    def on_train_epoch_end(self) -> None:
        if self.trainer.is_global_zero:
            self.train_embeddings = self.all_gather(self.train_embeddings).reshape(-1, self.train_embeddings.shape[-1])
            self.train_y = self.all_gather(self.train_y).reshape(-1)
            self.val_embeddings = self.all_gather(self.val_embeddings).reshape(-1, self.train_embeddings.shape[-1])
            self.val_y = self.all_gather(self.val_y).reshape(-1)
            # access all training embeddings and labels
            print('on_train_epoch_end! -- train embeddings: ', self.train_embeddings.shape,
                  ' -- val embeddings: ', self.val_embeddings.shape)
        else:
            print('on_train_epoch_end @!')

Since I want to perform KNN lookup, I added if self.trainer.is_global_zero control in order to get all knn related computation done one time and in master process.

But now the problem is training doesn't continue after on_train_epoch_end function executed.

When I check nvidia-smi it shows all 4 gpus utilized 100%.

It works when I remove if self.trainer.is_global_zero control but then all 4 processes do all computations.

Any recommendations?

[akihironitta]

Not sure if this will resolve your specific issue, but in general, if something that will be executed later depends on some rank-zero-specific logic, there should be barrier() to make sure other processes don't use information that hasn't been completed in the main process:

    def on_train_epoch_end(self) -> None:
        if self.trainer.is_global_zero:
            self.train_embeddings = self.all_gather(self.train_embeddings).reshape(-1, self.train_embeddings.shape[-1])
            self.train_y = self.all_gather(self.train_y).reshape(-1)
            self.val_embeddings = self.all_gather(self.val_embeddings).reshape(-1, self.train_embeddings.shape[-1])
            self.val_y = self.all_gather(self.val_y).reshape(-1)
            # access all training embeddings and labels
            print('on_train_epoch_end! -- train embeddings: ', self.train_embeddings.shape,
                  ' -- val embeddings: ', self.val_embeddings.shape)

+       self.trainer.strategy.barrier()  # other processes wait until the main process (the slowest process, precisely) reaches here

[samils7]

Thanks @akihironitta for this suggestion.

It doesn't work as is. The reason I found is self.all_gather() function takes some time and after main process calls first all_gather(), the others went into barrier(). Since all_gather() function needs other processes to be available in order to reach their resources, and somehow barrier() makes a process resources unreachable, later calls of all_gather() in master process ends up with deadlock.

I changed the structure of on_train_epoch_end func as:

def on_train_epoch_end(self) -> None:
    train_embeddings = self.all_gather(self.train_embeddings).reshape(-1, self.train_embeddings.shape[-1])
    train_y = self.all_gather(self.train_y).reshape(-1)
    val_embeddings = self.all_gather(self.val_embeddings).reshape(-1, self.train_embeddings.shape[-1])
    val_y = self.all_gather(self.val_y).reshape(-1)
    if self.trainer.is_global_zero:
        knn_acc = self.calculate_knn_accuracy(val_embeddings, train_embeddings, val_y, train_y, k=11)
        self.log('knn_acc', knn_acc, on_step=False, on_epoch=True, prog_bar=True, logger=True)

    self.trainer.strategy.barrier()

[akihironitta]

Since all_gather() function needs other processes to be available in order to reach their resources, and somehow barrier() makes a process resources unreachable, later calls of all_gather() in master process ends up with deadlock.

You're totally right! Thank you for coming back here and leaving the code snippet for the community.