Gradient flow about taming-transformers HOT 8 CLOSED

Thanks!

We use two different optimizers for the autoencoder part and the discriminator part, respectively. Furthermore, both optimizers have different optimiziation steps specified by the optimizer_idx (see taming.modules.losses.vqperceptual), so the discriminator should only receive updates when optimizer_idx = 1.

from taming-transformers.

I was asking myself the same thing. From what I understand from the lightning docs, they are called sequentially one after another in turn. At least unless users choose to overwrite this behaviour which I did not find in this repo.

from taming-transformers.

thanks for getting back to me on this one.

I understand the procedure you described, the optimizer index being owed to the fact that you use Lightning and all.

However, I am still convinced that the critic is receiving gradients in the backward pass before performing the step of the first optimizer since there is no disabling of its gradients anywhere before. That would mean that when its the critics turn to get its weights updated it would do this with gradients from both the backward pass through it all the way down to the VQVAE and from its own backward pass in the second stage of optimizing.

Unless I am missing sth. here (maybe lightning does something behind the curtain) I would find this rather peculiar since usually the critics weights are frozen when it is just judging the generators outcome...

from taming-transformers.

That is not correct, since the reconstructions are detached from the computation graph before passing them to the discriminator:

if cond is None:
    logits_real = self.discriminator(inputs.contiguous().detach())
    logits_fake = self.discriminator(reconstructions.contiguous().detach())
else:
    logits_real = self.discriminator(torch.cat((inputs.contiguous().detach(), cond), dim=1))
    logits_fake = self.discriminator(torch.cat((reconstructions.contiguous().detach(), cond), dim=1))

So the gradients just flow through the discriminator and then stop at the point where the inputs were detached, and no need to freeze the model.

(And of course, when training the vqVAE you cannot freeze the discriminator since you want the gradients to flow up to the model)

from taming-transformers.

@rromb Simple question,
I've understood that different optimizer steps are applied via "optimizer_idx". But, since I'm not used to PyTorch lightening, I am having trouble finding when the "optimization steps" are toggled. Can you tell me where the value of the "optimizer_idx" changes?

from taming-transformers.

@CDitzel Umm, I see. Thank you for your answer :)

from taming-transformers.

The code defining the two optimizers is in the VQGAN pl module. Just in case somone has the same issue.

from taming-transformers.

Pytorch-lighting dotoggle_optimizer and untoggle_optimizer to control which parameters to update. https://github.com/Lightning-AI/lightning/blob/1.8.0/src/pytorch_lightning/core/module.py#L1408-L1453

from taming-transformers.