Repository: pfnet/pytorch-pfn-extras

Issue:
Trainer: Support passing multiple evaluators



Discussion:

Resolving patch:
diff --git a/pytorch_pfn_extras/engine.py b/pytorch_pfn_extras/engine.py
--- a/pytorch_pfn_extras/engine.py
+++ b/pytorch_pfn_extras/engine.py
@@ -28,8 +28,10 @@ def create_trainer(
         out_dir: str = 'result',
         stop_trigger: 'TriggerLike' = None,
         writer: Optional['writing.Writer'] = None,
-        evaluator: Optional[
-            Union['Evaluator', Tuple['Evaluator', 'TriggerLike']]] = None,
+        evaluator: Optional[Union[
+            'Evaluator', Tuple['Evaluator', 'TriggerLike'],
+            Mapping[str, Union['Evaluator', Tuple['Evaluator', 'TriggerLike']]]
+        ]] = None,
         device: 'DeviceLike' = 'cpu',
         logic: Optional[handler_module.Logic] = None,
         transform_model: Callable[
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
@@ -1,3 +1,4 @@
+import collections.abc
 import queue
 import time
 import warnings
@@ -27,7 +28,9 @@ def __init__(
             self,
             handler: 'handler_module.BaseHandler',
             *,
-            evaluator: Optional[Union['Evaluator', Tuple['Evaluator', TriggerLike]]],
+            evaluator: Optional[Union[
+                'Evaluator', Tuple['Evaluator', TriggerLike],
+                Mapping[str, Union['Evaluator', Tuple['Evaluator', TriggerLike]]]]],
             models: Union[torch.nn.Module, Mapping[str, torch.nn.Module]],
             **kwargs: Any,
     ):
@@ -49,13 +52,14 @@ 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._evaluators: Dict[str, Tuple['Evaluator', TriggerLike]] = {}
+        if evaluator is None:
+            evaluator = {}
+        elif not isinstance(evaluator, collections.abc.Mapping):
+            evaluator = {"Evaluator": evaluator}
+        if isinstance(evaluator, collections.abc.Mapping):
+            for n, e in evaluator.items():
+                self._evaluators[n] = e if isinstance(e, tuple) else (e, (1, 'epoch'))
         self.val_loader = None
 
     def extend(
@@ -132,6 +136,14 @@ def stop_trigger(self) -> Trigger:
     def stop_trigger(self, trigger: Trigger) -> None:
         self._stop_trigger = trigger
 
+    @property
+    def evaluator(self) -> Optional['Evaluator']:
+        if len(self._evaluators) == 0:
+            return None
+        if len(self._evaluators) == 1:
+            return next(iter(self._evaluators.values()))[0]
+        raise ValueError('multiple evaluators are registered.')
+
     def get_optimizer(self, name: str) -> torch.optim.Optimizer:
         return self.manager.optimizers[name]
 
@@ -186,15 +198,7 @@ def _complete_step(
             record_run_iteration.complete()
             record_iteration.complete()
 
-    def _run_evaluator(self) -> None:
-        assert self.evaluator is not None
-        if self._val_loader is None:
-            raise ValueError('"val_loader" is not given.')
-        self.evaluator.handler.train_validation_begin(self, self.evaluator)
-        self.evaluator.run(self._val_loader, eval_len=self._eval_len)
-        self.evaluator.handler.train_validation_end(self, self.evaluator)
-
-    def run(self,  # type: ignore[override]
+    def run(self,
             train_loader: Iterable[Any],
             val_loader: Optional[Iterable[Any]] = None,
             *,
@@ -223,30 +227,43 @@ def run(self,  # type: ignore[override]
             eval_len = len(val_loader)  # type: ignore[arg-type]
 
         self._train_len = train_len
-        self._val_loader = val_loader
         self._eval_len = eval_len
 
         class _EvaluatorExt:
-            def __init__(self, trainer: 'Trainer') -> None:
-                self.name = 'Evaluator'
+            def __init__(
+                    self,
+                    trainer: 'Trainer',
+                    evaluator: 'Evaluator',
+                    val_loader: Optional[Iterable[Any]],
+                    eval_len: Optional[int],
+            ) -> None:
                 self.needs_model_state = True
                 self._trainer = trainer
+                self._evaluator = evaluator
+                self._val_loader = val_loader
+                self._eval_len = eval_len
 
             def __call__(self, manager: ExtensionsManagerProtocol) -> None:
-                self._trainer._run_evaluator()
+                evaluator = self._evaluator
+                if self._val_loader is None:
+                    raise ValueError('"val_loader" is not given.')
+                evaluator.handler.train_validation_begin(self._trainer, evaluator)
+                evaluator.run(self._val_loader, eval_len=self._eval_len)
+                evaluator.handler.train_validation_end(self._trainer, evaluator)
 
         if self._manager is None:
             self._manager = self._setup_manager(train_len)
-            if self.evaluator is not None:
+            for name, (evaluator, trigger) in self._evaluators.items():
                 # Register the evaluator as an extension to the manager
                 # To be triggered with the correct timing
                 self._manager.extend(
-                    _EvaluatorExt(self),
-                    trigger=self.evaluator_trigger,
+                    _EvaluatorExt(self, evaluator, val_loader, eval_len),
+                    name=name,
+                    trigger=trigger_module.get_trigger(trigger),
                     priority=extension.PRIORITY_WRITER,
                 )
             self.handler.train_setup(self, train_loader)
-            if self.evaluator is None:
+            if len(self._evaluators) == 0:
                 if val_loader is not None:
                     warnings.warn(
                         '`val_loader` is given whereas the evaluator is missing.',
@@ -254,7 +271,8 @@ def __call__(self, manager: ExtensionsManagerProtocol) -> None:
             else:
                 if val_loader is None:
                     raise ValueError('`val_loader` is required')
-                self.evaluator.handler.eval_setup(self.evaluator, val_loader)
+                for _, (evaluator, _) in self._evaluators.items():
+                    evaluator.handler.eval_setup(evaluator, val_loader)
 
         while not self.manager.stop_trigger:
             self.handler.train_epoch_begin(self, train_loader)


Test patch:
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
@@ -174,6 +174,36 @@ def test_evaluator_trigger(evaluator_trigger, path):
         assert patched.call_count == evaluator_trigger[0]
 
 
+def test_evaluator_dict(path):
+    device = 'cpu'
+    progress_bar = False
+    model = MyModel()
+    ppe.to(model, device)
+    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(10)])
+    extensions = _make_extensions()
+
+    evaluator1 = engine.create_evaluator(
+        model_with_loss, device=device, progress_bar=progress_bar)
+    evaluator2 = engine.create_evaluator(
+        model, device=device, progress_bar=progress_bar)
+
+    trainer = engine.create_trainer(
+        model_with_loss, optimizer, 20,
+        device=device, evaluator={
+            '1': evaluator1,  # called 20 times.
+            '2': (evaluator2, (5, 'iteration')),  # called 40 times.
+        },
+        extensions=extensions, out_dir=path
+    )
+    path = 'pytorch_pfn_extras.training._evaluator.Evaluator.run'
+    with mock.patch(path) as patched:
+        trainer.run(data, data)
+        assert patched.call_count == 20 + 40
+
+
 @pytest.mark.parametrize('device', ['cpu', 'cuda'])
 def test_train_result_equal(device, path):
     train_data = torch.utils.data.DataLoader(