Logging Images during validation using Tensorboard Logger

[muaali]

Hi,

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.

[rohitgr7]

hey @muaali !

by end of the validation round, I randomly select few images to log

you can log the images randomly inside validation_step itself to avoid accumulation that is creating memory overhead.

[Minxiangliu]

Hi @rohitgr7,
Do you have any suggestions on how to make tensorboard display images?

[akihironitta]

@Minxiangliu There's a callback in Bolts (although it's deprecated). You can use it as is or copy/modify the implementation as you need. Hope it helps. (Note that some components in Bolts may not work as it's not actively maintained.)
https://github.com/Lightning-AI/lightning-bolts/blob/master/pl_bolts/callbacks/vision/image_generation.py

[Minxiangliu]

@akihironitta Thank you.