Repository: pfnet/pytorch-pfn-extras

Issue:
Merge `_Engine` with `Trainer`
Now that only class inheriting `ppe.engine._Engine` is `ppe.training.Trainer`, we can merge these classes.


Discussion:

Resolving patch:
diff --git a/pytorch_pfn_extras/training/_trainer.py b/pytorch_pfn_extras/training/_trainer.py
--- a/pytorch_pfn_extras/training/_trainer.py
+++ b/pytorch_pfn_extras/training/_trainer.py
@@ -22,13 +22,15 @@
     from pytorch_pfn_extras.profiler._time_summary import _ReportNotification
 
 
-class _Engine:
+class Trainer:
     def __init__(
             self,
             handler: 'handler_module.BaseHandler',
+            *,
+            evaluator: Optional[Union['Evaluator', Tuple['Evaluator', TriggerLike]]],
             models: Union[torch.nn.Module, Dict[str, torch.nn.Module]],
             **kwargs: Any,
-    ) -> None:
+    ):
         self.handler = handler
         self._manager: Optional['training.ExtensionsManager'] = None
 
@@ -47,6 +49,15 @@ def __init__(
                   Dict[str, Any]]] = []
         self._manager_state: Optional[Dict[str, Any]] = None
 
+        if isinstance(evaluator, tuple):
+            self.evaluator: Optional['Evaluator'] = None
+            self.evaluator, trigger = evaluator
+            self.evaluator_trigger = trigger_module.get_trigger(trigger)
+        else:
+            self.evaluator = evaluator
+            self.evaluator_trigger = trigger_module.get_trigger((1, 'epoch'))
+        self.val_loader = None
+
     def extend(
             self,
             extension: 'training.Extension',
@@ -97,28 +108,6 @@ def load_state_dict(self, to_load: Dict[str, Any]) -> None:
             return
         self.manager.load_state_dict(to_load)
 
-    def run(self, *args: Any, **kwargs: Any) -> None:
-        raise NotImplementedError
-
-
-class Trainer(_Engine):
-    def __init__(
-            self,
-            handler: 'handler_module.BaseHandler',
-            *,
-            evaluator: Optional[Union['Evaluator', Tuple['Evaluator', TriggerLike]]],
-            **kwargs: Any,
-    ):
-        super().__init__(handler, **kwargs)
-        if isinstance(evaluator, tuple):
-            self.evaluator: Optional['Evaluator'] = None
-            self.evaluator, trigger = evaluator
-            self.evaluator_trigger = trigger_module.get_trigger(trigger)
-        else:
-            self.evaluator = evaluator
-            self.evaluator_trigger = trigger_module.get_trigger((1, 'epoch'))
-        self.val_loader = None
-
     @property
     def epoch(self) -> int:
         return self.manager.epoch


Test patch:
diff --git a/tests/pytorch_pfn_extras_tests/training_tests/test_engine.py b/tests/pytorch_pfn_extras_tests/training_tests/test_engine.py
--- a/tests/pytorch_pfn_extras_tests/training_tests/test_engine.py
+++ b/tests/pytorch_pfn_extras_tests/training_tests/test_engine.py
@@ -4,8 +4,8 @@
 
 class TestEngine:
     def test_engine_extension(self):
-        engine = ppe.training._trainer._Engine(
-            None, {}, optimizers={}, max_epochs=10)
+        engine = ppe.training._trainer.Trainer(
+            None, evaluator=None, models={}, optimizers={}, max_epochs=10)
         extension = ppe.training.extensions.LogReport()
         engine.extend(extension)
         # Create the actual manager object
@@ -16,15 +16,16 @@ def test_engine_extension(self):
     def test_engine_state_dict(self):
         manager = ppe.training.ExtensionsManager(
             {}, {}, 10, iters_per_epoch=100)
-        engine = ppe.training._trainer._Engine(
-            None, {}, optimizers={}, max_epochs=10)
+        engine = ppe.training._trainer.Trainer(
+            None, evaluator=None, models={}, optimizers={}, max_epochs=10)
         engine._setup_manager(100)
         assert engine.state_dict() == manager.state_dict()
 
     def test_engine_load_state_dict(self):
         manager = ppe.training.ExtensionsManager(
             {}, {}, 10, iters_per_epoch=100)
-        engine = ppe.training._trainer._Engine(None, {}, optimizers={}, max_epochs=1)
+        engine = ppe.training._trainer.Trainer(
+            None, evaluator=None, models={}, optimizers={}, max_epochs=1)
         engine.load_state_dict(manager.state_dict())
         engine._setup_manager(20)
         assert engine.state_dict() == manager.state_dict()
@@ -32,7 +33,8 @@ def test_engine_load_state_dict(self):
     def test_engine_load_state_dict_2(self):
         manager = ppe.training.ExtensionsManager(
             {}, {}, 10, iters_per_epoch=100)
-        engine = ppe.training._trainer._Engine(None, {}, optimizers={}, max_epochs=1)
+        engine = ppe.training._trainer.Trainer(
+            None, evaluator=None, models={}, optimizers={}, max_epochs=1)
         engine._setup_manager(20)
         engine.load_state_dict(manager.state_dict())
         assert engine.state_dict() == manager.state_dict()
@@ -41,18 +43,20 @@ def test_engine_load_state_dict_2(self):
 class TestEngineInvalid:
     def test_engine_wrong_models(self):
         with pytest.raises(ValueError, match='model must be an instance'):
-            ppe.training._trainer._Engine(None, object(), optimizers={}, max_epochs=10)
+            ppe.training._trainer.Trainer(
+                None, evaluator=None, models=object(), optimizers={}, max_epochs=10)
 
     def test_engine_not_started(self):
-        engine = ppe.training._trainer._Engine(None, {}, optimizers={}, max_epochs=10)
+        engine = ppe.training._trainer.Trainer(
+            None, evaluator=None, models={}, optimizers={}, max_epochs=10)
         with pytest.raises(RuntimeError, match='is not started'):
             engine.state_dict()
         with pytest.raises(RuntimeError, match='is not started'):
             engine.manager
 
     def test_extend_after_init(self):
-        engine = ppe.training._trainer._Engine(
-            None, {}, optimizers={}, max_epochs=10)
+        engine = ppe.training._trainer.Trainer(
+            None, evaluator=None, models={}, optimizers={}, max_epochs=10)
         engine._setup_manager(10)
         extension = ppe.training.extensions.LogReport()
         with pytest.raises(RuntimeError, match='cannot extend after'):