Visualization callbacks

[Parskatt]

The setup us is as follows:

I want to have a modular setup where I can in some cases visualize segmentation output of my network during training. I do not wish to pollute my lightningmodule with visualization code (e.g. using a logger in the training step), so I want to use a callback. I would like to both visualize the input images aswell as my outputs.

My options seems to be (maybe I've misunderstood), to either put the callback "on_train_batch_start" or "on_train_batch_end". on_train_batch_start makes the visualization of the input easy (but for obvious reasons cant show my output), but on_train_batch_end means (as I understand it) that I would have to attach my entire segmentation mask to my outputs, and I think this will end up devouring my memory if I keep saving these outputs during an entire epoch.

Is there a better way of doing this that I have not considered (without resorting to putting the visualization code directly in my training_step)?