Repository: pfnet/pytorch-pfn-extras

Issue:
Option argument is not applied during Logic initialization
It is possible to specify an `option` argument at the time of Logic class initialization.
https://github.com/pfnet/pytorch-pfn-extras/blob/3f023d33254a4744ea37fd5ea29724cfbb029d36/pytorch_pfn_extras/handler/_logic.py#L161

On the other hand, even if the options argument is not specified at the time of trainer initialization, it is treated as an empty dict.
https://github.com/pfnet/pytorch-pfn-extras/blob/3f023d33254a4744ea37fd5ea29724cfbb029d36/pytorch_pfn_extras/engine.py#L150

`Logic.cosume_options()` will update with default values for the empty dict values passed, so the options specified at initialization will no longer apply.
https://github.com/pfnet/pytorch-pfn-extras/blob/3f023d33254a4744ea37fd5ea29724cfbb029d36/pytorch_pfn_extras/handler/_logic.py#L185-L192

Consideration should be given to prohibiting `option` arguments at the time of initialization.


Discussion:

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
@@ -1,7 +1,18 @@
 import contextlib
 import dataclasses
 import warnings
-from typing import Any, Dict, Generator, Iterable, Mapping, Optional
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    Generator,
+    Iterable,
+    List,
+    Mapping,
+    Optional,
+    Tuple,
+    Union,
+)
 
 import torch
 from pytorch_pfn_extras.handler._code_block import forward, update_parameters
@@ -179,22 +190,45 @@ def __init__(
                 * ``'grad_scaler'`` (torch.cuda.amp.GradScaler):
                     A gradient scaler that outputs are applied to.
         """
-        super().__init__(options)
+        self.backward_outputs: Optional[
+            Union[str, List[str], Tuple[str, ...]]
+        ] = None
+        self._grad_scaler: Optional[torch.cuda.amp.GradScaler] = None
+        self._backward_fn: Optional[Callable[..., Any]] = None
+        self._autocast_options: Optional[Union[Dict[str, Any], bool]] = None
         self.model_name = model_name
 
+        super().__init__(options)
+
     def consume_options(self, options: Dict[str, Any]) -> None:
         super().consume_options(options)
 
-        self.backward_outputs = options.pop("backward_outputs", None)
-        self._grad_scaler = options.pop("grad_scaler", None)
+        self.backward_outputs = (
+            options.pop("backward_outputs", None) or self.backward_outputs
+        )
+        self._grad_scaler = (
+            options.pop("grad_scaler", None) or self._grad_scaler
+        )
+        self._backward_fn = (
+            options.pop("backward_function", None) or self._backward_fn
+        )
 
-        self._backward_fn = options.pop("backward_function", None)
-        autocast_options = options.pop("autocast", False)
-        if isinstance(autocast_options, bool):
+        # The `autocast` option is checked for `None`
+        # since it can take a bool value containing false.
+        autocast_options_value: Optional[
+            Union[Dict[str, Any], bool]
+        ] = options.pop("autocast", None)
+        if autocast_options_value is not None:
+            self._autocast_options = autocast_options_value
+
+        autocast_options_value = self._autocast_options or False
+        if isinstance(autocast_options_value, bool):
             autocast_options = {
-                "enabled": autocast_options,
+                "enabled": autocast_options_value,
                 "device_type": "cuda",
             }
+        else:
+            autocast_options = autocast_options_value
         self._autocast = _autocast._AutocastManager(
             autocast_options, self._grad_scaler is not None
         )
@@ -253,7 +287,7 @@ def _backward(self, outputs: Dict[str, Any]) -> None:
             assert (
                 len(to_backward) == 1
             ), "loss scaling with multiple loss is not supported"
-            to_backward = {self._grad_scaler.scale(v) for v in to_backward}
+            to_backward = {self._grad_scaler.scale(v) for v in to_backward}  # type: ignore[no-untyped-call]
         for v in to_backward:
             if self._backward_fn is None:
                 v.backward()  # type: ignore[no-untyped-call]
@@ -332,8 +366,8 @@ def train_step_optimizers(
         """
         optimizer = optimizers[self.model_name]
         if self._grad_scaler is not None:
-            self._grad_scaler.step(optimizer)
-            self._grad_scaler.update()
+            self._grad_scaler.step(optimizer)  # type: ignore[no-untyped-call]
+            self._grad_scaler.update()  # type: ignore[no-untyped-call]
         else:
             optimizer.step()
 


Test patch:
diff --git a/tests/pytorch_pfn_extras_tests/handler_tests/test_logic.py b/tests/pytorch_pfn_extras_tests/handler_tests/test_logic.py
--- a/tests/pytorch_pfn_extras_tests/handler_tests/test_logic.py
+++ b/tests/pytorch_pfn_extras_tests/handler_tests/test_logic.py
@@ -1,11 +1,13 @@
 import pathlib
-from typing import Any, Mapping
+from copy import deepcopy
+from typing import Any, Mapping, Tuple
 from unittest import mock
 
 import pytest
 import pytorch_pfn_extras as ppe
 import torch
 from torch import nn
+from torch.cuda.amp import GradScaler
 from torch.nn import Module
 from torch.nn import functional as F
 from torch.optim import Optimizer
@@ -140,3 +142,117 @@ def train_step(
     )
     trainer.run(data, data)
     assert backward_fn.call_count == epochs * iters_per_epoch
+
+
+logic_options_type_map = {
+    "backward_outputs": str,
+    "grad_scaler": GradScaler,
+    "backward_function": None,
+    "autocast": bool,
+}
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize(
+    "logic_options_name, trainer_options_name",
+    [
+        ((), ("backward_outputs", "backward_function", "autocast")),
+        (("backward_function",), ("backward_function", "autocast")),
+        ((), ("backward_outputs", "grad_scaler")),
+        (("backward_outputs", "backward_function"), ("backward_outputs",)),
+        (("backward_function", "autocast"), ("grad_scaler", "autocast")),
+        (
+            ("backward_outputs",),
+            ("backward_outputs", "grad_scaler", "backward_function"),
+        ),
+        (
+            ("backward_outputs", "grad_scaler", "autocast"),
+            ("backward_function",),
+        ),
+        (("autocast",), ("grad_scaler",)),
+        (("grad_scaler", "backward_function"), ()),
+        (("backward_function",), ("grad_scaler", "autocast")),
+    ],
+)
+def test_initialize_logic_with_options_with_autocast(
+    logic_options_name: Tuple[str, ...], trainer_options_name: Tuple[str, ...]
+):
+    logic_options = {
+        k: mock.MagicMock(spec=logic_options_type_map[k])
+        for k in logic_options_name
+    }
+    trainer_options = {
+        k: mock.MagicMock(spec=logic_options_type_map[k])
+        for k in trainer_options_name
+    }
+    expected_options = deepcopy(logic_options)
+    expected_options.update(deepcopy(trainer_options))
+    expected_logic = ppe.handler.Logic(options=expected_options)
+
+    actual_logic = ppe.handler.Logic(options=logic_options)
+    _ = ppe.engine.create_trainer(
+        models=mock.MagicMock(spec=nn.Module),
+        optimizers=mock.MagicMock(spec=Optimizer),
+        max_epochs=1,
+        logic=actual_logic,
+        options=trainer_options,
+    )
+    assert actual_logic.backward_outputs == expected_logic.backward_outputs
+    assert actual_logic._grad_scaler == expected_logic._grad_scaler
+    assert actual_logic._backward_fn == expected_logic._backward_fn
+    assert actual_logic._autocast._options == expected_logic._autocast._options
+
+
+# The `autocast` and `grad_scaler` option is removed from the test case
+# because it will result in a RuntimeError if the gpu is not present.
+@pytest.mark.parametrize(
+    "logic_options_name, trainer_options_name",
+    [
+        ((), ()),
+        ((), ("backward_outputs",)),
+        ((), ("backward_function",)),
+        ((), ("backward_outputs", "backward_function")),
+        (("backward_outputs",), ()),
+        (("backward_outputs",), ("backward_outputs",)),
+        (("backward_outputs",), ("backward_function",)),
+        (("backward_outputs",), ("backward_outputs", "backward_function")),
+        (("backward_function",), ()),
+        (("backward_function",), ("backward_outputs",)),
+        (("backward_function",), ("backward_function",)),
+        (("backward_function",), ("backward_outputs", "backward_function")),
+        (("backward_outputs", "backward_function"), ()),
+        (("backward_outputs", "backward_function"), ("backward_outputs",)),
+        (("backward_outputs", "backward_function"), ("backward_function",)),
+        (
+            ("backward_outputs", "backward_function"),
+            ("backward_outputs", "backward_function"),
+        ),
+    ],
+)
+def test_initialize_logic_with_options(
+    logic_options_name: Tuple[str, ...], trainer_options_name: Tuple[str, ...]
+):
+    logic_options = {
+        k: mock.MagicMock(spec=logic_options_type_map[k])
+        for k in logic_options_name
+    }
+    trainer_options = {
+        k: mock.MagicMock(spec=logic_options_type_map[k])
+        for k in trainer_options_name
+    }
+    expected_options = deepcopy(logic_options)
+    expected_options.update(deepcopy(trainer_options))
+    expected_logic = ppe.handler.Logic(options=expected_options)
+
+    actual_logic = ppe.handler.Logic(options=logic_options)
+    _ = ppe.engine.create_trainer(
+        models=mock.MagicMock(spec=nn.Module),
+        optimizers=mock.MagicMock(spec=Optimizer),
+        max_epochs=1,
+        logic=actual_logic,
+        options=trainer_options,
+    )
+    assert actual_logic.backward_outputs == expected_logic.backward_outputs
+    assert actual_logic._grad_scaler == expected_logic._grad_scaler
+    assert actual_logic._backward_fn == expected_logic._backward_fn
+    assert actual_logic._autocast._options == expected_logic._autocast._options
diff --git a/tests/pytorch_pfn_extras_tests/training_tests/test_trainer.py b/tests/pytorch_pfn_extras_tests/training_tests/test_trainer.py
--- a/tests/pytorch_pfn_extras_tests/training_tests/test_trainer.py
+++ b/tests/pytorch_pfn_extras_tests/training_tests/test_trainer.py
@@ -880,7 +880,7 @@ def test_trainer_with_clousure_logic(device, progress_bar, path):
         evaluator=evaluator,
         extensions=extensions,
         out_dir=path,
-        logic=ppe.handler.ClousureLogic(options={"backward_outputs": ["loss"]}),
+        logic=ppe.handler.ClousureLogic(),
     )
     trainer.run(data, data)