Repository: pfnet/pytorch-pfn-extras
Issue: Recovering the mode of the model after running evaluator

When the evaluator is run, `ppe.handler.logic.train_validation_begin` changes the model to `eval` mode, but the `ppe.handler.logic.train_ validation_end`, which is executed after the evaluator is run, does not restore it.

Conventionally, it is specified in `train` mode via `ppe.handler.logic.train_epoch_begin`, but if the timing of the evaluator execution is at a non-integer epoch frequency, it is possible that a training that is not in `train` mode will be executed.

Therefore, the mode of the model is restored at `ppe.handler.logic.train_validation_end`.

https://github.com/pfnet/pytorch-pfn-extras/blob/bcb5f37eba7ffd48f052231c886df21ab759a0f8/pytorch_pfn_extras/handler/_logic.py#L330-L340


Changed files:
tests/pytorch_pfn_extras_tests/test_logic.py, pytorch_pfn_extras/handler/_logic.py

diff --git a/pytorch_pfn_extras/handler/_logic.py b/pytorch_pfn_extras/handler/_logic.py
index b13e1fd..b9dacc6 100644
--- a/pytorch_pfn_extras/handler/_logic.py
+++ b/pytorch_pfn_extras/handler/_logic.py
@@ -268,6 +268,10 @@ class Logic(BaseLogic):
             # Needed for `torch.utils.data.DistributedSampler`
             loader.sampler.set_epoch(epoch)  # type: ignore[attr-defined]
 
+    def train_epoch_end(self, models: Mapping[str, Any], epoch: int) -> None:
+        model = models[self.model_name]
+        model.eval()
+
     def train_step(
             self,
             models: Mapping[str, torch.nn.Module],
@@ -339,6 +343,10 @@ class Logic(BaseLogic):
         model = models[self.model_name]
         model.eval()
 
+    def train_validation_end(self, models: Mapping[str, Any]) -> None:
+        model = models[self.model_name]
+        model.train()
+
     def eval_step(
             self,
             models: Mapping[str, torch.nn.Module],
@@ -404,6 +412,10 @@ class CodeBlockLogic(BaseLogic):
             # Needed for `torch.utils.data.DistributedSampler`
             loader.sampler.set_epoch(epoch)  # type: ignore[attr-defined]
 
+    def train_epoch_end(self, models: Mapping[str, Any], epoch: int) -> None:
+        model = models[self.model_name]
+        model.eval()
+
     def train_step(
             self,
             models: Mapping[str, torch.nn.Module],
@@ -447,6 +459,10 @@ class CodeBlockLogic(BaseLogic):
         model = models[self.model_name]
         model.eval()
 
+    def train_validation_end(self, models: Mapping[str, Any]) -> None:
+        model = models[self.model_name]
+        model.train()
+
     def eval_step(
             self,
             models: Mapping[str, torch.nn.Module],
diff --git a/tests/pytorch_pfn_extras_tests/test_logic.py b/tests/pytorch_pfn_extras_tests/test_logic.py
index eade4d6..9d85f88 100644
--- a/tests/pytorch_pfn_extras_tests/test_logic.py
+++ b/tests/pytorch_pfn_extras_tests/test_logic.py
@@ -1,8 +1,10 @@
+from typing import Any, Mapping
 import pytest
 
 import torch
 from torch import nn
-from torch.nn import functional as F
+from torch.optim import Optimizer
+from torch.nn import functional as F, Module
 from unittest import mock
 
 import pytorch_pfn_extras as ppe
@@ -60,3 +62,39 @@ def test_trainer(device):
     )
     trainer.run(data)
     assert backward_fn.call_count == epochs * iters_per_epoch
+
+
+@pytest.mark.parametrize("trigger", [(1, "epoch"), (0.5, "epoch"), (10, "iteration"), (5, "iteration"), (1, "iteration")])
+def test_train_step_mode_with_evaluator(trigger):
+    iters_per_epoch = 10
+    epochs = 20
+    model = MyModel()
+    ppe.to(model, "cpu")
+    model_with_loss = MyModelWithLossFn(model)
+    optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
+    data = torch.utils.data.DataLoader(
+        [(torch.rand(20,), torch.rand(10,)) for i in range(iters_per_epoch)])
+    backward_fn = mock.Mock(return_value=None)
+
+    class LogicWithTrainStepCheck(ppe.handler.Logic):
+        def train_step(self, models: Mapping[str, Module], optimizers: Mapping[str, Optimizer], batch_idx: int, batch: Any) -> Any:
+            model = models[self.model_name]
+            assert model.training
+            return super().train_step(models, optimizers, batch_idx, batch)
+
+    trainer = ppe.engine.create_trainer(
+        model_with_loss,
+        optimizer,
+        epochs,
+        logic=LogicWithTrainStepCheck(),
+        evaluator=(
+            ppe.engine.create_evaluator(
+                models=model_with_loss,
+                logic=LogicWithTrainStepCheck(),
+            ),
+            trigger,
+        ),
+        options={'backward_function': backward_fn}
+    )
+    trainer.run(data, data)
+    assert backward_fn.call_count == epochs * iters_per_epoch