Isn't here duplicate computation for gan example?

[wztdream]

Hi,
It seems there is duplicate computation for gan example.
One here:

# train generator
        if optimizer_idx == 0:

            # generate images
            self.generated_imgs = self(z) ## ，<-- first

The other here:

# train discriminator
        if optimizer_idx == 1:
....

fake_loss = self.adversarial_loss(self.discriminator(self(z).detach()), fake) <--- repeat

And also the sampling seems also duplicate:

def training_step(self, batch, batch_idx, optimizer_idx):
        imgs, _ = batch

        # sample noise
        z = torch.randn(imgs.shape[0], self.hparams.latent_dim)
        z = z.type_as(imgs)

I do think this is unreasonable.

So how to avoid this unnecessary computation in lightning, I know there are some tricks to achieve, such as using a self.some_state to hold the values and the reuse them. But they seems no elegant.

[akihironitta]

@wztdream Yes, it is redundant. I see a few options:

1. Store the computed result (the noise and generated images) in some attribute and reuse it for the next optimizer indexed optimizer_idx==1.
2. Use manual optimization.
   • Docs: https://pytorch-lightning.readthedocs.io/en/stable/model/manual_optimization.html?highlight=manual%20optimization#manual-optimization
   • Example: https://github.com/akihironitta/gist/blob/5f29d7aa8712e2f16c929b1384b555dd2e53ae9f/pl_gan_manual_optimization/main.py
3. (Use the loop API. There's an example of yielding training step using the loop API: https://github.com/Lightning-AI/lightning/blob/master/examples/pl_loops/yielding_training_step.py, but I'd recommend the first two above because this one needs a bit more work.)

[wztdream]

thank you, I am new to pytorch lightning, although these "workaround" are good enough to use, but I still have feeling that training_step may need redesign to better handle these situations.

[akihironitta]

Thanks for your feedback :) I think that's exactly what manual optimization is made for.