Training discriminator/generator pair in GAN

[MandaloreUltimate]

I'm currently trying to build DCGAN for MNIST and got stuck on the simultaneous training of discriminator and generator as I'm really confused on what should be included in the respective loss functions. The first code snippet results in zero gradients, then I tried moving the calculation of 'generated_images', 'real_output' and 'fake_output' under the loss functions (second snippet) and it technically started training but generator was producing weird images, sort of collapsing.

The only example I could find to refer to was this pull request, but it tries to utilize single loss function and as it's stated in the comment, it's not going to work. It's also noted that there should be a single optimizer for the model, but how is it going to handle both discriminator and generator at the same time? The only way I could come up with is to train discriminator independently, fix its new weights and stack it with generator into new 'GAN' model, but that's not what I need.

@jax.jit
def train_step(generator_optimizer, discriminator_optimizer, images):
    noise = jax.random.normal(random.PRNGKey(0), shape = [256, 100])
    generated_images = generator_optimizer.target(noise)
    real_output = discriminator_optimizer.target(images)
    fake_output = discriminator_optimizer.target(generated_images)
    
    def generator_loss(model):
        return binary_cross_entropy(jnp.ones_like(fake_output), fake_output).mean(), generated_images
    
    def discriminator_loss(model):
        real_loss = binary_cross_entropy(jnp.ones_like(real_output), real_output).mean()
        fake_loss = binary_cross_entropy(jnp.zeros_like(fake_output), fake_output).mean()
        total_loss = real_loss + fake_loss
        return total_loss, fake_output

    grad_fn_gen = jax.value_and_grad(generator_loss, has_aux=True)
    (_, preds), grad_gen = grad_fn_gen(generator_optimizer.target)
    generator_optimizer = generator_optimizer.apply_gradient(grad_gen)
    
    grad_fn_disc = jax.value_and_grad(discriminator_loss, has_aux=True)
    (_, preds), grad_disc = grad_fn_disc(discriminator_optimizer.target)
    discriminator_optimizer = discriminator_optimizer.apply_gradient(grad_disc)
    
    return generator_optimizer, discriminator_optimizer

def generator_loss(generator):
        generated_images = generator(noise)
        real_output = discriminator_optimizer.target(images)
        fake_output = discriminator_optimizer.target(generated_images)
        return binary_cross_entropy(jnp.ones_like(fake_output), fake_output).mean(), generated_images
    
    def discriminator_loss(discriminator):
        generated_images = generator_optimizer.target(noise)
        real_output = discriminator(images)
        fake_output = discriminator(generated_images)
    
        real_loss = binary_cross_entropy(jnp.ones_like(real_output), real_output).mean()
        fake_loss = binary_cross_entropy(jnp.zeros_like(fake_output), fake_output).mean()
        total_loss = real_loss + fake_loss
        return total_loss, fake_output

[jheek]

Because generator and discriminator have different loss functions we need 2 separate calls to jax.value_and_grad. The second snippet is more correct but it still has an issue: You are using optimizer.target instead of the model passed to the loss function. This way gradients aren't computed so what you need is the following:

def generator_loss(generator):
    generated_images = generator(noise)
        
    def discriminator_loss(discriminator):
        real_output = discriminator(images)
        fake_output = discriminator(generated_images)
    
        real_loss = binary_cross_entropy(jnp.ones_like(real_output), real_output).mean()
        fake_loss = binary_cross_entropy(jnp.zeros_like(fake_output), fake_output).mean()
        total_loss = real_loss + fake_loss
        return total_loss, fake_output
  
  (disc_loss, fake_output), disc_grad = jax.value_and_grad(discriminator_loss)(discriminator_optimizer.target)
  gen_loss = binary_cross_entropy(jnp.ones_like(fake_output), fake_output).mean()
  return gen_loss, (generated_images, disc_loss, disk_grad)