Remove training_step returned None user warning when automatic_optimization=False

[timothybrooks]

🐛 Bug

When using manual optimization by setting the property automatic_optimization=False, it is not necessary (and sometimes undesirable) for training_step() to return a loss output. For example, in the case of GANs or other models with complex multi-optimizer setups, manual optimization can be preferable or necessary for the correct behavior. The user warning should therefore be omitted in this case.

Since the user calls the update step themselves (self.manual_backward(); optimizer.step()) it is not necessary for the training step to return a loss since PL does not need to use the loss to update weights. Furthermore, logging the returned output loss is unhelpful for certain complex multi-optimizer setups, since aggregation of losses from different optimizers is not desired.

In short, since the updates and logging when using automatic_optimization=False is non-standard and does not always involve returning a loss output, I believe the UserWarning("Your training_step returned None. Did you forget to return an output?") should be omitted in the case of manual optimization.

Basic example for GAN:

@property
def automatic_optimization(self) -> bool:
  return False

def training_step(...):
  loss_d = self.step_and_optimize_d(...)
  self.log("loss_d", loss_d)

  loss_g = self.step_and_optimize_g(...)
  self.log("loss_g", loss_g)

  # Note that the training step does not return a value.

[carmocca]

Is this not the case already in master?

https://github.com/PyTorchLightning/pytorch-lightning/blob/6166f46281d77830cf88b1cb8dc050c0312fcfcf/pytorch_lightning/trainer/training_loop.py#L743-L765

what version are you using?

[timothybrooks]

I am using the stable release 1.2, so if it has since been addressed in master that is great. In my version this block is nearly the same, but line 743 uses self.trainer.train_loop.automatic_optimization instead of self.automatic_optimization. Perhaps this was the cause of the issue since automatic_optimization migrated from a training flag to a module property? In any case, if the warnings are omitted correctly on master during manual optimization than it is no longer an issue.

[rohitgr7]

@timothybrooks are you training in distributed env (ddp, dp, ...)??

[timothybrooks]

Yes in DDP.

…

On Fri, Mar 5, 2021 at 6:42 AM Rohit Gupta ***@***.***> wrote: @timothybrooks <https://github.com/timothybrooks> are you training in distributed env (ddp, dp, ...)?? — You are receiving this because you were mentioned. Reply to this email directly, view it on GitHub <#6339 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/ACQMGH2BENRLOO7X63PUO3TTCDUTPANCNFSM4YSFFZWA> .

[timothybrooks]

Hello wanted to check in since I am still experiencing this bug in 1.2.6 when automatic_optimization=False and the training step returns none. Since the bug has been closed for a few weeks I wanted to check that it is still planned to be addressed. Thanks!

[akihironitta]

@timothybrooks I think we somehow missed the fix #6139. It should be included in the next release. It seems the fix is blocked by #4945, so it should be included in 1.3.

[timothybrooks]

Got it, thanks for the update. Looking forward to it!