Why my parameters do not learn/update?

[oschan77]

Hello there, I am training a GAN model (StyleGAN2-ADA) and I find that the parameters are not updating at all. I am using quite a lot of custom layers and blocks. I wonder if it causes the parameters not learning at all. How can I make my model learn and solve this issue? Thank you so much!

my training_step function:

  def training_step(self, batch, batch_idx, optimizer_idx):
    ### Get a batch of real images and possibly apply augmentation
    real_images = self.augmenter(batch, training=True)

    ### Get w and noise
    self.w = self.get_w()
    self.noise = self.get_noise()

    ## Get current epoch
    ep = self.current_epoch + 1

    ## Discriminator
    if optimizer_idx == 0:
      ### Check whether to apply lazy regularization
      do_lazy_reg = (ep % self.lazy_gradient_penalty_interval) == 0

      for i in range(self.gradient_accumulate_steps):
        ### Generate images by the generator
        generated_images, _ = self()
        ### Possibly apply augmentation on the generated images
        generated_images = self.augmenter(generated_images, training=True)
        ### Put the generated images into the discriminator
        fake_output = self.discriminator(generated_images.detach())

        if (do_lazy_reg):
          ### Gradients w.r.t. real images are calculated for gradient penalty
          real_images.requires_grad_()

        ### Put the real images into the discriminator
        real_output = self.discriminator(real_images)

        ### Update the augmentation probability based on the performance of discriminator 
        self.augmenter.update(real_output)

        ### Get the discriminator loss (WGAN's loss)
        real_loss, fake_loss = self.discriminator_loss(real_output, fake_output)
        disc_loss = real_loss + fake_loss

        if (do_lazy_reg):
          ### Calculate the gradient penalty term
          gp = self.gradient_penalty(real_images, real_output)
          ### Calculate the new discriminator loss by adding the gradient penalty term
          disc_loss = disc_loss + 0.5 * self.gradient_penalty_coefficient * gp * self.lazy_gradient_penalty_interval

      return {'disc_loss': disc_loss}

    ## Generator
    if (optimizer_idx == 1) or (optimizer_idx == 2):
      ### Check whether to apply path length penalty (plp)
      do_plp = (ep >= self.lazy_path_penalty_after) and (ep % self.lazy_path_penalty_interval == 0)

      for i in range(self.gradient_accumulate_steps):
        ### Generate images by the generator
        generated_images, w = self()
        ### Put the generated images into the discriminator
        fake_output = self.discriminator(generated_images)
        ### Get the generator loss (WGAN's loss)
        gen_loss = self.generator_loss(fake_output)

        if (do_plp):
          ### Calculate the path length penalty (plp) term
          plp = self.path_length_penalty(w, generated_images)
          if not torch.isnan(plp):
            ### Calculate the new generator loss by adding the path length penalty (plp) term
            gen_loss = gen_loss + plp

      return {'gen_loss': gen_loss}

[rohitgr7]

return value from training_step must have a 'loss' key.

return {'loss': gen_loss}

reF: https://pytorch-lightning.readthedocs.io/en/stable/common/lightning_module.html#training-step