GenerateImage#

class dmme.callbacks.GenerateImage(imgsize, timesteps, batch_size=8, vis_length=20, every_n_epochs=5)[source]#

Generate samples to check training progress

Parameters:

  • imgsize (Tuple[int, int, int]) – A tuple of ints representing image shape \((C, H, W)\)

  • batch_size (int) – Number of samples to generate

  • vis_length (int) – Length of denoising sequence to visualize

  • every_n_epochs (int) – Only save those images every N epochs

  • test (bool) – generates images on test if set to true

on_train_epoch_end(trainer, pl_module)[source]#

Called when the train epoch ends.

To access all batch outputs at the end of the epoch, either:

  1. Implement training_epoch_end in the LightningModule and access outputs via the module OR

  2. Cache data across train batch hooks inside the callback implementation to post-process in this hook.