What's the best way to do two generator steps?

Trainer
#1

Dear Pytorch lightning community,
I am using Pytorch lightning to train a GAN. Thus, for each training_step I have one generator_step and one discriminator_step (similar to https://github.com/PyTorchLightning/pytorch-lightning-bolts/blob/f48357be353b7acdd882379ac3308fbec95dc40d/pl_bolts/models/gans/basic/basic_gan_module.py#L119).

I need to do two steps of the discriminator_step, thus I have to backward two times for each batch. What’s the best way to do it in pytorch lighting?

thank you

#2

you mean for a single batch you want to do:

# Generator
gen_loss.backward()
gen_opt.step()
gen_opt.zero_grad()

# Discriminator
disc_loss.backward()
disc_opt.step()
disc_opt.step()
disc_opt.zero_grad()

or something else?

#3

I would like to do this

Generator

gen_loss.backward()
gen_opt.step()
gen_opt.zero_grad()

Discriminator

disc_loss.backward()
disc_opt.step()
disc_opt.zero_grad()
disc_loss2.backward()
disc_opt.step()
disc_opt.zero_grad()

#4

See manual optimization. This should allow you to do exactly what you want!

1 Like
#5

I think you might get an error here if you use the old weights to calculate disc_loss2 because first disc_opt.step() will update the weights and disc_loss2.backward() will be calculating the gradients using the new weights.

1 Like
#6

@goku thank you, it was a wrong example indeed but @teddy solution works!! Thank you all