remove backward from training batch loop (#9265)

Author: awaelchli

CHANGELOG.md

@@ -67,6 +67,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Loop customization:
     * Added `Closure` and `AbstractClosure` classes ([#8642](https://github.com/PyTorchLightning/pytorch-lightning/pull/8642))
     * Refactored `TrainingBatchLoop` and extracted `OptimizerLoop`, splitting off automatic optimization into its own loop ([#9191](https://github.com/PyTorchLightning/pytorch-lightning/pull/9191))
+    * Removed `TrainingBatchLoop.backward()`; manual optimization now calls directly into `Accelerator.backward()` and automatic optimization handles backward in new `OptimizerLoop` ([#9265](https://github.com/PyTorchLightning/pytorch-lightning/pull/9265))
 
 - Added support for saving and loading state of multiple callbacks of the same type ([#7187](https://github.com/PyTorchLightning/pytorch-lightning/pull/7187))
 

pytorch_lightning/core/lightning.py

@@ -1418,7 +1418,7 @@ def training_step(...):
         self._verify_is_manual_optimization("manual_backward")
 
         # backward
-        self.trainer.fit_loop.epoch_loop.batch_loop.backward(loss, None, None, *args, **kwargs)
+        self.trainer.accelerator.backward(loss, None, None, *args, **kwargs)
 
     def backward(
         self, loss: Tensor, optimizer: Optional[Optimizer], optimizer_idx: Optional[int], *args, **kwargs

pytorch_lightning/loops/batch/training_batch_loop.py

@@ -16,7 +16,6 @@
 from typing import Any, Callable, List, Optional, Tuple
 
 import numpy as np
-import torch
 from deprecate import void
 from torch import Tensor
 from torch.optim import Optimizer
@@ -246,25 +245,6 @@ def _tbptt_split_batch(self, batch: Any) -> List[Any]:
             splits = model_ref.tbptt_split_batch(batch, tbptt_steps)
         return splits
 
-    # TODO: remove this method and update tests
-    def backward(
-        self,
-        loss: Tensor,
-        optimizer: Optional[torch.optim.Optimizer],
-        opt_idx: Optional[int] = None,
-        *args: Any,
-        **kwargs: Any,
-    ) -> Tensor:
-        """Performs the backward step.
-
-        Args:
-            loss: The loss value to back-propagate on
-            optimizer: Current optimizer being used. ``None`` if using manual optimization.
-            opt_idx: Index of the current optimizer being used. ``None`` if using manual optimization.
-        """
-        self.trainer.accelerator.backward(loss, optimizer, opt_idx, *args, **kwargs)
-        return loss
-
     def _update_running_loss(self, current_loss: Tensor) -> None:
         """Updates the running loss value with the current value"""
         if self.trainer.lightning_module.automatic_optimization:

tests/trainer/test_trainer.py

@@ -910,7 +910,7 @@ def test_gradient_clipping_by_norm(tmpdir, precision):
         gradient_clip_val=1.0,
     )
 
-    old_backward = trainer.fit_loop.epoch_loop.batch_loop.backward
+    old_backward = trainer.fit_loop.epoch_loop.batch_loop.optimizer_loop.backward
 
     def backward(*args, **kwargs):
         # test that gradient is clipped correctly
@@ -920,7 +920,7 @@ def backward(*args, **kwargs):
         assert (grad_norm - 1.0).abs() < 0.01, f"Gradient norm != 1.0: {grad_norm}"
         return ret_val
 
-    trainer.fit_loop.epoch_loop.batch_loop.backward = backward
+    trainer.fit_loop.epoch_loop.batch_loop.optimizer_loop.backward = backward
     trainer.fit(model)
 
 
@@ -945,7 +945,7 @@ def test_gradient_clipping_by_value(tmpdir, precision):
         default_root_dir=tmpdir,
     )
 
-    old_backward = trainer.fit_loop.epoch_loop.batch_loop.backward
+    old_backward = trainer.fit_loop.epoch_loop.batch_loop.optimizer_loop.backward
 
     def backward(*args, **kwargs):
         # test that gradient is clipped correctly
@@ -958,7 +958,7 @@ def backward(*args, **kwargs):
         ), f"Gradient max value {grad_max} != grad_clip_val {grad_clip_val} ."
         return ret_val
 
-    trainer.fit_loop.epoch_loop.batch_loop.backward = backward
+    trainer.fit_loop.epoch_loop.batch_loop.optimizer_loop.backward = backward
     trainer.fit(model)