Repository: pfnet/pytorch-pfn-extras

Issue:
Support passing `closure` argument to optimizers
Optimizers supports two kind of `step()` call:

https://pytorch.org/docs/stable/optim.html#taking-an-optimization-step

However, some optimizers require mandatory `closure` argument:

https://pytorch.org/docs/stable/generated/torch.optim.LBFGS.html#torch.optim.LBFGS

Currently, ExtensionsManager and Logic do not support these optimizers.



Discussion:
Since `torch.cuda.amp.GradScaler` does not support the step method by closuer, it must be processed by branching.

https://pytorch.org/docs/stable/amp.html#torch.cuda.amp.GradScaler.step

The `torch.optim.obfgs.LBFGS` class, which requires a clousure, assumes that the clousure returns an object that can be converted to a float.

https://pytorch.org/docs/stable/_modules/torch/optim/lbfgs.html#LBFGS.step

Resolving patch:
diff --git a/pytorch_pfn_extras/handler/__init__.py b/pytorch_pfn_extras/handler/__init__.py
--- a/pytorch_pfn_extras/handler/__init__.py
+++ b/pytorch_pfn_extras/handler/__init__.py
@@ -1,6 +1,6 @@
 from pytorch_pfn_extras.handler._code_block import CodeBlock, update_parameters, forward  # NOQA
 from pytorch_pfn_extras.handler._handler import BaseHandler, Handler  # NOQA
-from pytorch_pfn_extras.handler._logic import BaseLogic, Logic, CodeBlockLogic  # NOQA
+from pytorch_pfn_extras.handler._logic import BaseLogic, Logic, CodeBlockLogic, ClousureLogic  # NOQA
 
 # Deprecated, only imported for backward compatibility
 from pytorch_pfn_extras.handler._logic import torch_autocast  # NOQA
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,4 +1,5 @@
 import contextlib
+import dataclasses
 from typing import Any, Dict, Generator, Iterable, Mapping, Optional
 import warnings
 
@@ -464,3 +465,80 @@ def eval_step(
         model = models[self.model_name]
         outs = forward(model)(batch)
         return outs
+
+
+@dataclasses.dataclass
+class ClousureModelOutput:
+    outs: Any
+    loss: torch.Tensor
+
+    def __float__(self) -> float:
+        return float(self.loss)
+
+
+class ClousureLogic(Logic):
+
+    def consume_options(self, options: Dict[str, Any]) -> None:
+        super().consume_options(options)
+        if self._grad_scaler is not None:
+            raise RuntimeError('torch.cuda.amp.GradScaler does not support clousure step mode.')
+
+    def train_step(
+            self,
+            models: Mapping[str, torch.nn.Module],
+            optimizers: Mapping[str, torch.optim.Optimizer],
+            batch_idx: int,
+            batch: Any,
+    ) -> Any:
+        """A method invokes the model forward and backward passes and performs an optimization step.
+
+        Args:
+            models (dict of torch.nn.Module):
+                The models.
+            optimizers (dict of torch.optim.Optimizer):
+                The optimizers.
+            batch_idx (int):
+                Number of training steps already finished.
+            batch (torch.Tensor, list of torch.Tensor, dict of torch.Tensor):
+                Input tensors feeded to the model of the current step.
+        """
+        def clousure() -> ClousureModelOutput:
+            with torch_autocast(enabled=self._autocast):
+                optimizers[self.model_name].zero_grad()
+                outs = self._forward(models[self.model_name], batch)
+            to_back_outs = _normalize_outputs(outs)
+            if len(to_back_outs) > 1:
+                raise RuntimeError("Clousure step with multiple outputs is not supported.")
+            elif len(to_back_outs) == 0:
+                raise RuntimeError("No backward target found.")
+
+            self._backward(to_back_outs)
+            loss, = to_back_outs.values()
+            return ClousureModelOutput(
+                outs=outs,
+                loss=loss,
+            )
+
+        optimizer = optimizers[self.model_name]
+        clousure_model_output: ClousureModelOutput = optimizer.step(clousure)  # type: ignore
+        if not isinstance(clousure_model_output, ClousureModelOutput):
+            raise RuntimeError(f"{type(clousure_model_output)} type object returned from optimizer.step with clousure. optimizer.step is expected to return ppe.handler.ClousureModelOutput.")
+        return clousure_model_output.outs
+
+    def train_step_optimizers(
+            self,
+            models: Mapping[str, torch.nn.Module],
+            optimizers: Mapping[str, torch.optim.Optimizer],
+            batch_idx: int,
+    ) -> None:
+        """In clousure mode, the stepping of the optimizer cannot be changed.
+
+        If you want to change the stepping of the optimizer, please use the normal Logic class.
+
+        Args:
+            optimizers (dict of torch.optim.Optimizer):
+                The optimizers.
+            batch_idx (int):
+                Number of steps already finished.
+        """
+        pass


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
@@ -529,3 +529,9 @@ def test_eval_step(self):
         models['main'].eval()
         out = logic.eval_step(models, 0, input)
         torch_testing_assert_close(out, model(input))
+
+    @pytest.mark.gpu
+    def test_use_grad_scaler_with_clousure(self):
+        options = {'grad_scaler': torch.cuda.amp.GradScaler()}
+        with pytest.raises(RuntimeError):
+            ppe.handler.ClousureLogic(options=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
@@ -506,3 +506,28 @@ def test_trainer_profile():
     )
     trainer.run(data, data)
     assert trace_handler.call_count == 20  # n_epochs
+
+
+@pytest.mark.parametrize('device', ['cpu', 'cuda'])
+@pytest.mark.parametrize('progress_bar', [True, False])
+def test_trainer_with_clousure_logic(device, progress_bar, path):
+    if not torch.cuda.is_available() and device == 'cuda':
+        pytest.skip()
+    model = MyModel()
+    model_with_loss = MyModelWithLossFn(model)
+    ppe.to(model_with_loss, device)
+    optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
+    data = torch.utils.data.DataLoader(
+        [{'x': torch.rand(20,), 't': torch.rand(10,)} for i in range(10)])
+    extensions = _make_extensions()
+
+    evaluator = engine.create_evaluator(
+        model_with_loss, device=device, progress_bar=progress_bar,
+        logic=ppe.handler.ClousureLogic())
+
+    trainer = engine.create_trainer(
+        model_with_loss, optimizer, 20,
+        device=device, evaluator=evaluator, extensions=extensions,
+        out_dir=path, logic=ppe.handler.ClousureLogic(options={"backward_outputs": ["loss"]})
+    )
+    trainer.run(data, data)