Repository: Project-MONAI/MONAI

Issue:
Besides `default_prepare_batch`, add support to prepare extra data from batch
**Is your feature request related to a problem? Please describe.**
Currently, we have a basic `default_prepare_batch` in engines, it would be useful to add support to prepare extra data besides `image` and `label` as our engines already support `args` and `kwargs`:
https://github.com/Project-MONAI/MONAI/blob/dev/monai/engines/trainer.py#L181
We have reference implementation in NVIDIA Clara, will reimplement it in MONAI stype.



Discussion:

Resolving patch:
diff --git a/monai/engines/__init__.py b/monai/engines/__init__.py
--- a/monai/engines/__init__.py
+++ b/monai/engines/__init__.py
@@ -15,6 +15,9 @@
 from .utils import (
     GanKeys,
     IterationEvents,
+    PrepareBatch,
+    PrepareBatchDefault,
+    PrepareBatchExtraInput,
     default_make_latent,
     default_metric_cmp_fn,
     default_prepare_batch,
diff --git a/monai/engines/utils.py b/monai/engines/utils.py
--- a/monai/engines/utils.py
+++ b/monai/engines/utils.py
@@ -9,13 +9,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Sequence, Tuple, Union
 
 import torch
 
 from monai.config import IgniteInfo
 from monai.transforms import apply_transform
-from monai.utils import min_version, optional_import
+from monai.utils import ensure_tuple, min_version, optional_import
 from monai.utils.enums import CommonKeys
 
 if TYPE_CHECKING:
@@ -28,6 +29,9 @@
     "GanKeys",
     "get_devices_spec",
     "default_prepare_batch",
+    "PrepareBatch",
+    "PrepareBatchDefault",
+    "PrepareBatchExtraInput",
     "default_make_latent",
     "engine_apply_transform",
     "default_metric_cmp_fn",
@@ -123,6 +127,80 @@ def default_prepare_batch(
     return batchdata[CommonKeys.IMAGE].to(device=device, non_blocking=non_blocking), None
 
 
+class PrepareBatch(ABC):
+    """
+    Interface of customized prepare_batch in the trainer or evaluator workflows.
+    It takes the data of current batch, target device and non_blocking flag as input.
+
+    """
+
+    @abstractmethod
+    def __call__(
+        self,
+        batchdata: Dict[str, torch.Tensor],
+        device: Optional[Union[str, torch.device]] = None,
+        non_blocking: bool = False,
+    ):
+        raise NotImplementedError(f"Subclass {self.__class__.__name__} must implement this method.")
+
+
+class PrepareBatchDefault(PrepareBatch):
+    """
+    Default prepare_batch method to return `image` and `label` only,
+    it's consistent with `default_prerpare_batch` API.
+
+    """
+
+    def __call__(
+        self,
+        batchdata: Dict[str, torch.Tensor],
+        device: Optional[Union[str, torch.device]] = None,
+        non_blocking: bool = False,
+    ):
+        return default_prepare_batch(batchdata, device, non_blocking)
+
+
+class PrepareBatchExtraInput(PrepareBatch):
+    """
+    Customized prepare_batch for trainer or evalutor that support extra input data for network.
+    Extra items are specified by the `extra_keys` parameter.
+
+    Args:
+        extra_keys: if a string or list provided, every item is the key of extra data in current batch,
+            and will pass the extra data to the network(*args) in order.
+            if a dict provided, every `{k, v}` pair is the key of extra data in current batch,
+            `k` the param name in network, `v` is the key of extra data in current batch,
+            and will pass the `{k1: batch[v1], k2: batch[v2], ...}` as kwargs to the network.
+
+    """
+
+    def __init__(self, extra_keys: Union[str, Sequence[str], Dict[str, str]]) -> None:
+        self.extra_keys = extra_keys
+
+    def __call__(
+        self,
+        batchdata: Dict[str, torch.Tensor],
+        device: Optional[Union[str, torch.device]] = None,
+        non_blocking: bool = False,
+    ):
+        image, label = default_prepare_batch(batchdata, device, non_blocking)
+        args = list()
+        kwargs = dict()
+
+        def _get_data(key: str):
+            data = batchdata[key]
+            return data.to(device=device, non_blocking=non_blocking) if isinstance(data, torch.Tensor) else data
+
+        if isinstance(self.extra_keys, (str, list, tuple)):
+            for k in ensure_tuple(self.extra_keys):
+                args.append(_get_data(k))
+        elif isinstance(self.extra_keys, dict):
+            for k, v in self.extra_keys.items():
+                kwargs.update({k: _get_data(v)})
+
+        return image, label, tuple(args), kwargs
+
+
 def default_make_latent(
     num_latents: int, latent_size: int, device: Optional[Union[str, torch.device]] = None, non_blocking: bool = False
 ) -> torch.Tensor:


Test patch:
diff --git a/tests/min_tests.py b/tests/min_tests.py
--- a/tests/min_tests.py
+++ b/tests/min_tests.py
@@ -144,6 +144,8 @@ def run_testsuit():
         "test_zoomd",
         "test_transchex",
         "test_handler_mlflow",
+        "test_prepare_batch_extra_input",
+        "test_prepare_batch_default",
     ]
     assert sorted(exclude_cases) == sorted(set(exclude_cases)), f"Duplicated items in {exclude_cases}"
 
diff --git a/tests/test_prepare_batch_default.py b/tests/test_prepare_batch_default.py
new file mode 100644
--- /dev/null
+++ b/tests/test_prepare_batch_default.py
@@ -0,0 +1,54 @@
+# Copyright 2020 - 2021 MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import torch
+
+from monai.engines import PrepareBatchDefault, SupervisedEvaluator
+from tests.utils import assert_allclose
+
+
+class TestNet(torch.nn.Module):
+    def forward(self, x: torch.Tensor):
+        return x
+
+
+class TestPrepareBatchDefault(unittest.TestCase):
+    def test_content(self):
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        dataloader = [
+            {
+                "image": torch.tensor([1, 2]),
+                "label": torch.tensor([3, 4]),
+                "extra1": torch.tensor([5, 6]),
+                "extra2": 16,
+                "extra3": "test",
+            }
+        ]
+        # set up engine
+        evaluator = SupervisedEvaluator(
+            device=device,
+            val_data_loader=dataloader,
+            epoch_length=1,
+            network=TestNet(),
+            non_blocking=False,
+            prepare_batch=PrepareBatchDefault(),
+            decollate=False,
+        )
+        evaluator.run()
+        output = evaluator.state.output
+        assert_allclose(output["image"], torch.tensor([1, 2], device=device))
+        assert_allclose(output["label"], torch.tensor([3, 4], device=device))
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/test_prepare_batch_extra_input.py b/tests/test_prepare_batch_extra_input.py
new file mode 100644
--- /dev/null
+++ b/tests/test_prepare_batch_extra_input.py
@@ -0,0 +1,76 @@
+# Copyright 2020 - 2021 MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import torch
+from parameterized import parameterized
+
+from monai.engines import PrepareBatchExtraInput, SupervisedEvaluator
+from tests.utils import assert_allclose
+
+TEST_CASE_0 = [
+    {"extra_keys": "extra1"},
+    {"x": torch.tensor([1, 2]), "t1": torch.tensor([5, 6]), "t2": None, "t3": None},
+]
+
+TEST_CASE_1 = [
+    {"extra_keys": ["extra1", "extra3"]},
+    {"x": torch.tensor([1, 2]), "t1": torch.tensor([5, 6]), "t2": "test", "t3": None},
+]
+
+TEST_CASE_2 = [
+    {"extra_keys": {"t1": "extra2", "t2": "extra3", "t3": "extra1"}},
+    {"x": torch.tensor([1, 2]), "t1": 16, "t2": "test", "t3": torch.tensor([5, 6])},
+]
+
+
+class TestNet(torch.nn.Module):
+    def forward(self, x: torch.Tensor, t1=None, t2=None, t3=None):
+        return {"x": x, "t1": t1, "t2": t2, "t3": t3}
+
+
+class TestPrepareBatchExtraInput(unittest.TestCase):
+    @parameterized.expand([TEST_CASE_0, TEST_CASE_1, TEST_CASE_2])
+    def test_content(self, input_args, expected_value):
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        dataloader = [
+            {
+                "image": torch.tensor([1, 2]),
+                "label": torch.tensor([3, 4]),
+                "extra1": torch.tensor([5, 6]),
+                "extra2": 16,
+                "extra3": "test",
+            }
+        ]
+        # set up engine
+        evaluator = SupervisedEvaluator(
+            device=device,
+            val_data_loader=dataloader,
+            epoch_length=1,
+            network=TestNet(),
+            non_blocking=True,
+            prepare_batch=PrepareBatchExtraInput(**input_args),
+            decollate=False,
+        )
+        evaluator.run()
+        output = evaluator.state.output
+        assert_allclose(output["image"], torch.tensor([1, 2], device=device))
+        assert_allclose(output["label"], torch.tensor([3, 4], device=device))
+        for k, v in output["pred"].items():
+            if isinstance(v, torch.Tensor):
+                assert_allclose(v, expected_value[k].to(device))
+            else:
+                self.assertEqual(v, expected_value[k])
+
+
+if __name__ == "__main__":
+    unittest.main()