Logging Images during validation using Tensorboard Logger


I have successfully implemented the method to log images to tensorboard logger, except I run out of GPU memory soon as I accumulate images during the whole validation_step and by end of the validation round, I randomly select few images to log. This is not the best way to do it. Can someone point me how to do it properly where I don’t consume too much of Memory. Thanks.

This is how my validation step looks:

def validation_step(self, batch, batch_idx):
        imgs, y_true = batch
        y_pred = self(imgs)
        val_loss = self.nn_criterion(y_pred, y_true)
        self.log("val_loss", val_loss)
        return {"val_loss": val_loss,
                "images": imgs,
                "masks_pred": y_pred,
                "true_masks": y_true}

One can clearly see that I am accumulating tensors over the validation step. Since I am working with very large dataset, I run out of memory very soon. Thanks in advance.

If I can somehow access a single batch of validation inside validation_epoch_end that should solve the problem actually.

you can save the batch inside the state variable and access it?

def validation_step(self, batch, batch_idx):
    self.val_batch = batch

def validation_epoch_end(self, outputs):
    self.val_batch # access it here

Also, we have moved the discussions to GitHub Discussions. You might want to check that out instead to get a quick response. The forums will be marked read-only after some time.

Thank you