GAN training with multiple optimizers: do I need to call generator.train() and generator.eval()?

[silviu-oprea]

Looking at the basic GAN example here: What if the generator or discriminator has dropout and batchnorm layers?

• When optimizing the discriminator, dropout and batchnorm should be disabled in the generator.
• In vanilla PyTorch, this can be achieved by calling generator.eval(). However, the example linked above, under training_step, under if optimizer_idx == 1, does not call generator.eval().
• Is this called somewhere under the hood? I.e. when the discriminator optimizer is used, is there an .eval() called on all other parameters (the ones not changed by the current optimizer, i.e. the generator parameters)?
• At the moment I am calling generator.eval() manually. Is this doing something wrong overall?

[awaelchli]

Lightning does not know if you are running a GAN nor does it know which part of the model is a generator and which one is a discriminator. When execution enters training_step, you can always assume that all layers are in train mode. If this is not desired then you will need to toggle it yourself as you explained, and that's perfectly fine to do. Never shy away of writing PyTorch code in PyTorch Lightning! :))

[awaelchli]

Also, you might want to consider doing that in one of the hooks so you don't do it every training_step.