Add nrow and normalize args to TensorboardGenerativeModelImageSampler

[chris-clem]

🚀 Feature

Add nrow and normalize args to TensorboardGenerativeModelImageSampler

Motivation

At the moment, the image grid that is added to Tensorboard is created with grid = torchvision.utils.make_grid(images). It would be nice to have control over the nrow and normalize args of make_grid, e.g. through grid = torchvision.utils.make_grid(images, nrow=self.nrow, normalize=self.normalize).

Pitch

Change TensorboardGenerativeModelImageSampler's init from

    def __init__(self, num_samples: int = 3):
        super().__init__()
        self.num_samples = num_samples

to

    def __init__(self, num_samples: int = 3, nrow: int = 1, normalize: bool = True):
        super().__init__()
        self.num_samples = num_samples
        self.nrow = nrow
        self.normalize = normalize

and grid = torchvision.utils.make_grid(images) to grid = torchvision.utils.make_grid(images, nrow=self.nrow, normalize=self.normalize).

Alternatives

• Add args for all of make_grid's args with sensible default values
• Create a general image logging callback that can deal with all loggers and not just Tensorboard

[akihironitta]

@chris-clem Thank you for your suggestion! +1 for this change :]

It seems torchvision.utils.make_grid has other arguments as well, which we could possibly wrap them here. Do you think we should add other parameters, too, like padding, pad_value, ... of the function?

cc: @Borda @PyTorchLightning/bolts-contributors

[chris-clem]

Yes, I think wrapping all parameters and using sensible defaults makes sense while we are already at it.

[carmocca]

Makes sense, we can just pass **kwargs with the same default values.

[akihironitta]

@chris-clem Would you be interested in submitting a PR?

[chris-clem]

@chris-clem Would you be interested in submitting a PR?

@akihironitta Yes, I will do it in the next days.