Repository: Lightning-AI/pytorch-lightning

Issue:
GAN training with Pytorch Lightning is broken.
I was trying to train a DCGAN on my dataset but it wouldn't work in any means until I detach the training logic from Lightning and run the code without it. It was not working when my training logic is in Lightning module. I checked the gan examples in the docs and also multiple optimizer things. After 2 days of headaches, source code inspections and putting numerous print statements in the lightning source code, I found the culprit.

GAN training with Pytorch Lightning is simply broken. The culprit is only calling optimizer.zero_grad() after optimizer.step() since it clears the gradients of Generator or Discriminator only. Before the other network's weights are updated for, say, Generator; "loss.backward()" is called and it updates gradients for all parameters, but when optimizer.zero_grad() is called after the parameters are updated, only Generator's gradients are cleaned. So when it comes the Discriminator loss.backward(), leftover gradients are accumulated for Discriminator parameters and it just messes up with everything. Any kind of GAN training is impossible with this settings. That's why you can not find any GAN implementations with Pytorch Lightning on the internet.

Possible solutions
-Putting a warning in the docs or on the console after detecting multiple optimizers are defined like "if you train GAN, don't forget to zero all gradients by overriding the optimizer_step method or just reset gradients in your training loop before returning loss dictionary". (would be weird, honestly)

- calling zero_grad before calling backward in optimizer_closure (but it would mess with gradient accumulation, i suppose)

- calling self.zero_grad() instead of optimizer.zero_grad() by default in the optimizer_step.

- call zero_grad for all optimizers of the lightning module after gradient step by default so that the hook on_before_zero_grad is called whenever appropriate.

(By the way, on_before_zero_grad is not seem to be called anywhere right now. Maybe the issue can be fixed with a default behavior of that method, alternatively.)


I just sent a pull request which implements the last option and updates the docs.










Discussion:
sorry for the trouble! good catch though. Thanks for the PR, we’ll merge it once tests pass :)

feel free to submit the PR when ready
Whoops, I just have failed to submit the PR. I think it is ready now. Appreciate your work very much. Thanks for the great library :) 
hi, at work we train GANs successfully without any optimizer overrides. We also have the zero_grad before optimizer step

my teammate recommends that you put no_grad around your discriminator forward, and your runtime would be faster than calling zero_grad in the proposed way. Would that work for you? 

Hi, I'm that teammate. Strictly speaking you shouldn't even need to zero gradients at the start of optimizer_step, or have to overload optimizer_step. Between `with torch.no_grad():` and `Tensor.detach()`, you can prevent unwanted gradient accumulation.
Hi, teammates.
 I was getting good results as I added a zero_grad before my training loop but the problem is, the training loop should not include these tricks for the generic training regimes like GAN in the first place. The mentality of Lightning is to just define a loss with respect to your models and leave that kind of boilerplate to the framework, that's why I thought it would be better to submit a PR instead of hacking my training loop or advising it in the docs as I mentioned in the PR message :) 
can you guys post pseudocode for both versions to see what you mean more clearly?

```
def training_step(self, batch, batch_nb, optimizer_idx):
        self.zero_grad()
        x_data = batch
        batch_size = x_data.shape[0]
        z = torch.randn(batch_size, 100).cuda()
        x_gen = self.G(z)
        real_labels = torch.ones(batch_size).cuda() 
        fake_labels = torch.zeros(batch_size).cuda() 
        
        if optimizer_idx == 0:           
            
            D_real = self.D(x_data)
            D_fake = self.D(x_gen)
            
            real_loss = self.loss(D_real, real_labels)
            fake_loss = self.loss(D_fake, fake_labels)
            d_loss = real_loss + fake_loss

            return {
                'loss': d_loss,
                'progress_bar': {'d_loss': d_loss}
                }
  
    
        if optimizer_idx == 1:

            D_fake = self.D(x_gen)
            g_loss = self.loss(D_fake, real_labels)

            return {
                'loss': g_loss,
                'progress_bar': {'g_loss': g_loss}
                }
```


this code works but it fails without self.zero_grad()
Yes, it would. Make the following changes to fix it:

Prevent gradients accumulating on the generator during the discriminator pass:
`D_fake = self.D(x_gen)` -> `D_fake = self.D(x_gen.detach())`

Prevent gradients accumulating on the discriminator during the generator pass:
`D_fake = self.D(x_gen)` ->
```
with torch.no_grad():
    D_fake = self.D(x_gen)
```

Hope that helps.
I've tried these and it works. I mean don't get me wrong, it's a really smart and elegant way to handle things, it even speeds up the code by ~10% but I think people shouldn't think about what gradients to prevent and what to let when training different models when it can be solved by a simple change at the high level that covers most of the cases. I really can't think of a case why one needs to hold any gradients after optimizer step. Also, your no_grad method would still be relevant after this PR.
Ok, sounds like we actually want to make sure users think about this when training GANs. But totally agree that it’s a bit annoying. 

Let’s make these changes:
1. no change to the code
2. update docs to help users understand this. 

If it turns out that other users run into this as well, we can make it the default. 

thank you all for a great discussion!

Resolving patch:
diff --git a/pytorch_lightning/core/lightning.py b/pytorch_lightning/core/lightning.py
--- a/pytorch_lightning/core/lightning.py
+++ b/pytorch_lightning/core/lightning.py
@@ -664,6 +664,8 @@ def configure_optimizers(self):
         .. note:: If you use multiple optimizers, training_step will have an additional `optimizer_idx` parameter.
 
         .. note:: If you use LBFGS lightning handles the closure function automatically for you.
+        
+        .. note:: If you use multiple optimizers, gradients will be calculated only for the parameters of current optimizer at each training step.
 
         Example
         -------
diff --git a/pytorch_lightning/trainer/training_loop.py b/pytorch_lightning/trainer/training_loop.py
--- a/pytorch_lightning/trainer/training_loop.py
+++ b/pytorch_lightning/trainer/training_loop.py
@@ -433,6 +433,13 @@ def run_training_batch(self, batch, batch_idx):
 
             # call training_step once per optimizer
             for opt_idx, optimizer in enumerate(self.optimizers):
+                # make sure only the gradients of the current optimizer's paramaters are calculated 
+                # in the training step to prevent dangling gradients in multiple-optimizer setup.
+                for param in self.get_model().parameters():
+                    param.requires_grad = False
+                for group in optimizer.param_groups:
+                    for param in group['params']:
+                        param.requires_grad = True
 
                 # wrap the forward step in a closure so second order methods work
                 def optimizer_closure():


Test patch: