Repository: pfnet/pytorch-pfn-extras

Issue:
Raise a warning when no backward happen
This is to give users a chance to notice an error (typo etc) in `backward_outputs` option.

https://github.com/pfnet/pytorch-pfn-extras/blob/a7dd62145bdd65cd1aa6e8a14c9c6c83ca160f11/pytorch_pfn_extras/handler/_logic.py#L215-L233


Discussion:
Why not raise an error when specified backward outputs are not found in `target`?

Resolving patch:
diff --git a/pytorch_pfn_extras/handler/_logic.py b/pytorch_pfn_extras/handler/_logic.py
--- a/pytorch_pfn_extras/handler/_logic.py
+++ b/pytorch_pfn_extras/handler/_logic.py
@@ -2,6 +2,8 @@
 from typing import (
     Any, Dict, Generator, Iterable, Optional,
 )
+import warnings
+
 from pytorch_pfn_extras.handler._code_block import forward, update_parameters
 
 _amp_enabled = False
@@ -215,22 +217,33 @@ def _forward(self, model: torch.nn.Module, batch: Any) -> Any:
 
     def _backward(self, outputs: Dict[str, Any]) -> None:
         target = _normalize_outputs(outputs)
-
-        for k, v in target.items():
-            # This is to avoid errors when the trained models returns
-            # tensors others than scalars
-            if isinstance(v, torch.Tensor) and v.grad_fn is not None and (
-                (
-                    self.backward_outputs is None
-                    and v.numel() == 1
-                    and (v.dtype.is_floating_point or v.dtype.is_complex)
-                )
-                or (
-                    self.backward_outputs is not None
-                    and k in self.backward_outputs
-                )
-            ):
-                v.backward()  # type: ignore[no-untyped-call]
+        to_backward = set()
+        if self.backward_outputs is None:
+            for _, v in target.items():
+                if isinstance(v, torch.Tensor) and v.grad_fn is not None and (
+                    (
+                        v.numel() == 1
+                        and (v.dtype.is_floating_point or v.dtype.is_complex)
+                    )
+                ):
+                    to_backward.add(v)
+        else:
+            # If backward is requested, we tried to execute it no matter the
+            # shape or type of the tensor to make the user aware
+            backward_outputs = self.backward_outputs
+            if type(backward_outputs) is str:
+                backward_outputs = (backward_outputs,)
+            for k in backward_outputs:
+                try:
+                    to_backward.add(target[k])
+                except KeyError:
+                    warnings.warn(
+                        'Couldn\'t find requested backward value: '
+                        f'{k} in {target.keys()}'
+                    )
+
+        for v in to_backward:
+            v.backward()  # type: ignore[no-untyped-call]
 
     def train_epoch_begin(
             self,


Test patch:
diff --git a/tests/pytorch_pfn_extras_tests/test_handler.py b/tests/pytorch_pfn_extras_tests/test_handler.py
--- a/tests/pytorch_pfn_extras_tests/test_handler.py
+++ b/tests/pytorch_pfn_extras_tests/test_handler.py
@@ -550,6 +550,27 @@ def forward(self, x):
 
         assert outs['0'].grad is None
 
+    def test_train_step_backward_invalid(self):
+        logic = ppe.handler.Logic(options={'backward_outputs': 'abcd'})
+        input = torch.rand(1, 1)
+        input.requires_grad = True
+
+        class _DummyModel(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.l0 = torch.nn.Linear(1, 1)
+
+            def forward(self, x):
+                return {'0': x}
+
+        model = _DummyModel()
+        models = {'main': model}
+        optimizers = {'main': torch.optim.SGD(model.parameters(), 1.0)}
+        assert input.grad is None
+
+        with pytest.warns(UserWarning, match='backward value: abcd'):
+            logic.train_step(models, optimizers, 0, input)
+
     def test_train_step_optimizers(self):
         logic = ppe.handler.Logic()
         models, optimizers, input, out = self._run_step(logic, 'cpu')