Best Practices: logger.experiment.add_image() at end of epoch when using new simplified pl.Train/EvalResult objects
Question:
What are we going to consider best practice for visualizing images, embeddings, etc. to tensorboard when using pl.Train/EvalResult objects?
In light of #2651 and related PRs, what's the right way to do this?
Let's say we have a dataset of images, and we want to visualize a single batch of reconstructions once per epoch. I typically do this in validation_epoch_end(), using logger.experiment.add_image().
Code:
Let's say my code now looks like this:
 def validation_step(self, batch, batch_idx):
        step_dict = self._step(batch)
        result = pl.EvalResult(early_stop_on=step_dict['loss'],
                               checkpoint_on=step_dict['loss'],
                              )

        # logging
        result.log('avg_val_loss', step_dict['loss'],
                   on_epoch=True, reduce_fx=torch.mean)

        return result, step_dict

which works fine and is definitely much cleaner than the original method of returning multiple logging dicts. 😁
I now want to do something at the end of the validation loop so I specify:
def validation_epoch_end(self, outputs):
        img_batch = outputs[-1]['step_dict']['viz']['x_hat']
        img_batch = img_batch.view(img_batch.shape[0], *[1,32,32])
        grid = torchvision.utils.make_grid(img_batch)
        self.logger.experiment.add_image('x_hat', grid, self.current_epoch)

        avg_val_loss = torch.stack([x['avg_val_loss'] for x,y in outputs]).mean()
        return avg_val_loss

outputs in this case is a list of tuples where the first element is the EvalResult for each val step, and the second element contains step_dict which includes all losses and reconstructed x_hats for each val step.
Is there a better way? One potential downside to this is that outputs can eat up a significant chunk of memory if you're not careful.
What's your environment?

OS: [Ubuntu]
Packaging [pip]
Version [master]