How to save/load hyperparameters of system componenets?

[ahrzb]

How should I handle hyperparameters of a submodule? For example in the snippet below, say MyGenerator takes it's hparams in constructor, how should I handle that?

class LitMNIST(LightningModule):

    def __init__(self, loss_fx, generator_network, layer_1_dim=128 **kwargs):
        super().__init__()
        self.layer_1_dim = layer_1_dim
        self.loss_fx = loss_fx

        # call this to save (layer_1_dim=128) to the checkpoint
        self.save_hyperparameters('layer_1_dim')

# to load specify the other args
model = LitMNIST.load_from_checkpoint(PATH, loss_fx=torch.nn.SomeOtherLoss, generator_network=MyGenerator(???))

[tchaton]

Hey @ahrzb,

I would say you could try a different approach there. Have a look at our example: https://github.com/PyTorchLightning/pytorch-lightning/blob/master/pl_examples/domain_templates/generative_adversarial_net.py#L122

You shouldn't pass the generator as an argument, but only the arguments to build it.
Also I would advice to check out Hydra. It could be used to dynamically instantiate your Generator.

class GAN(LightningModule):
    """
    >>> GAN(img_shape=(1, 8, 8))  # doctest: +ELLIPSIS +NORMALIZE_WHITESPACE
    GAN(
      (generator): Generator(
        (model): Sequential(...)
      )
      (discriminator): Discriminator(
        (model): Sequential(...)
      )
    )
    """
    def __init__(
            self,
            img_shape: tuple = (1, 28, 28),
            lr: float = 0.0002,
            b1: float = 0.5,
            b2: float = 0.999,
            latent_dim: int = 100,
    ):
        super().__init__()

        self.save_hyperparameters()

        # networks
        self.generator = Generator(latent_dim=self.hparams.latent_dim, img_shape=img_shape)
        self.discriminator = Discriminator(img_shape=img_shape)

        self.validation_z = torch.randn(8, self.hparams.latent_dim)

        self.example_input_array = torch.zeros(2, self.hparams.latent_dim)

[davidcian]

This is also a bit of an issue for me. I have a model with three components (a convolutional one, one made of ConvLSTMs and a fully connected classifier). I have to try fairly diverse components for these three slots, so I'm not sure how I can instantiate them inside my wrapper LightningModule...

[ohayonguy]

@tchaton any suggestion on how to do that anyways? I want to save a nn.Module as a "hyper parameter". The solution you offered is not scalable for me. I have a lightning module with many different components that can change drastically, and I want to do this as generic as possible so as to control each module from the configuration file (I am using Lightning CLI). If I wouldn't be able to do that, then this will cause a lot of boilerplate in my case (consider text recognition models for instance).

Saving the modules as hyper parameters seems to cause some issues, but if we don't save them as hyper parameters, then model loading does not work of course.

What do you suggest?