Repository: pfnet/pytorch-pfn-extras

Issue:
Trainer: Support `NamedTuple` in addition to `dict` as an input/output



Discussion:

Resolving patch:
diff --git a/pytorch_pfn_extras/handler.py b/pytorch_pfn_extras/handler.py
--- a/pytorch_pfn_extras/handler.py
+++ b/pytorch_pfn_extras/handler.py
@@ -515,14 +515,18 @@ def set_options(self, options):
         self._grad_scaler = options.pop('grad_scaler', None)
 
     def _forward(self, model, batch):
+        if isinstance(batch, tuple) and hasattr(batch, '_fields'):
+            return model(batch)
         if isinstance(batch, dict):
             return model(**batch)
-        elif isinstance(batch, (list, tuple)):
+        if isinstance(batch, (list, tuple)):
             return model(*batch)
         return model(batch)
 
     def _normalize_outputs(self, outputs):
-        if isinstance(outputs, dict):
+        if isinstance(outputs, tuple) and hasattr(outputs, '_fields'):
+            target = {k: getattr(outputs, k) for k in outputs._fields}
+        elif isinstance(outputs, dict):
             target = outputs
         elif isinstance(outputs, (list, tuple)):
             target = {str(i): out for i, out in enumerate(outputs)}
diff --git a/pytorch_pfn_extras/runtime/_runtime.py b/pytorch_pfn_extras/runtime/_runtime.py
--- a/pytorch_pfn_extras/runtime/_runtime.py
+++ b/pytorch_pfn_extras/runtime/_runtime.py
@@ -36,17 +36,20 @@ def convert_batch(self, args):
 
         # this should be called with the runtime associated to a model
         # or a model part
+        if isinstance(args, tuple) and hasattr(args, '_fields'):
+            return args.__class__(
+                **{k: self.move_tensor(getattr(args, k)) for k in args._fields})
         if isinstance(args, dict):
             return {
                 k: self.move_tensor(v) if isinstance(v, torch.Tensor) else v
                 for k, v in args.items()
             }
-        elif isinstance(args, (list, tuple)):
+        if isinstance(args, (list, tuple)):
             return [
                 self.move_tensor(v) if isinstance(v, torch.Tensor) else v
                 for v in args
             ]
-        elif isinstance(args, torch.Tensor):
+        if isinstance(args, torch.Tensor):
             return self.move_tensor(args)
         return args
 
diff --git a/pytorch_pfn_extras/training/extensions/evaluator.py b/pytorch_pfn_extras/training/extensions/evaluator.py
--- a/pytorch_pfn_extras/training/extensions/evaluator.py
+++ b/pytorch_pfn_extras/training/extensions/evaluator.py
@@ -202,7 +202,9 @@ def evaluate(self):
                 progress.current_position = idx
                 observation = {}
                 with reporting.report_scope(observation):
-                    if isinstance(batch, (tuple, list)):
+                    if isinstance(batch, tuple) and hasattr(batch, '_fields'):
+                        outs = self.eval_func(batch)
+                    elif isinstance(batch, (tuple, list)):
                         outs = self.eval_func(*batch)
                     elif isinstance(batch, dict):
                         outs = self.eval_func(**batch)


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
@@ -1,4 +1,5 @@
 import tempfile
+import typing
 
 import pytest
 
@@ -364,3 +365,83 @@ def test_trainer_autoload(self, path):
         assert new_trainer.epoch == 20
         assert _compare_states(
             trainer.state_dict(), new_trainer.state_dict())
+
+
+class MyModelWithLossDictOutput(torch.nn.Module):
+    def __init__(self, model):
+        super().__init__()
+        self.model = model
+
+    def forward(self, x, t):
+        y = self.model(x)
+        prefix = 'train' if self.training else 'val'
+        loss = F.l1_loss(y, t)
+        ppe.reporting.report({prefix + '/loss': loss})
+        return {'y': y, 'loss': loss}
+
+
+@pytest.mark.parametrize('device', ['cpu', 'cuda'])
+@pytest.mark.parametrize('progress_bar', [True, False])
+def test_trainer_dict_input(device, progress_bar, path):
+    model = MyModel()
+    ppe.to(model, device)
+    model_with_loss = MyModelWithLossDictOutput(model)
+    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)
+
+    trainer = engine.create_trainer(
+        model_with_loss, optimizer, 20,
+        device=device, evaluator=evaluator, extensions=extensions,
+        out_dir=path
+    )
+    trainer.run(data, data)
+
+
+class Input(typing.NamedTuple):
+    x: torch.Tensor
+    t: torch.Tensor
+
+
+class Output(typing.NamedTuple):
+    y: torch.Tensor
+    loss: torch.Tensor
+
+
+class MyModelWithLossNamedTupleOutput(torch.nn.Module):
+    def __init__(self, model):
+        super().__init__()
+        self.model = model
+
+    def forward(self, input):
+        y = self.model(input.x)
+        prefix = 'train' if self.training else 'val'
+        loss = F.l1_loss(y, input.t)
+        ppe.reporting.report({prefix + '/loss': loss})
+        return Output(y, loss)
+
+
+@pytest.mark.parametrize('device', ['cpu', 'cuda'])
+@pytest.mark.parametrize('progress_bar', [True, False])
+def test_trainer_namedtuple_input(device, progress_bar, path):
+    model = MyModel()
+    ppe.to(model, device)
+    model_with_loss = MyModelWithLossNamedTupleOutput(model)
+    optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
+    data = torch.utils.data.DataLoader(
+        [Input(torch.rand(20,), torch.rand(10,)) for i in range(10)])
+    extensions = _make_extensions()
+
+    evaluator = engine.create_evaluator(
+        model_with_loss, device=device, progress_bar=progress_bar)
+
+    trainer = engine.create_trainer(
+        model_with_loss, optimizer, 20,
+        device=device, evaluator=evaluator, extensions=extensions,
+        out_dir=path
+    )
+    trainer.run(data, data)