enhance copy in metatensor (#4506)

* shallow copy in meta structure

Signed-off-by: Wenqi Li <[email protected]>

* is batch no deepcopy

Signed-off-by: Wenqi Li <[email protected]>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: wyli

monai/data/meta_obj.py

@@ -11,8 +11,9 @@
 
 from __future__ import annotations
 
+import itertools
 from copy import deepcopy
-from typing import Any, Callable, Sequence
+from typing import Any, Iterable
 
 from monai.utils.enums import TraceKeys
 
@@ -74,86 +75,88 @@ class MetaObj:
     """
 
     def __init__(self):
-        self._meta: dict = self.get_default_meta()
-        self._applied_operations: list = self.get_default_applied_operations()
+        self._meta: dict = MetaObj.get_default_meta()
+        self._applied_operations: list = MetaObj.get_default_applied_operations()
         self._is_batch: bool = False
 
     @staticmethod
-    def flatten_meta_objs(args: Sequence[Any]) -> list[MetaObj]:
+    def flatten_meta_objs(*args: Iterable):
         """
-        Recursively flatten input and return all instances of `MetaObj` as a single
-        list. This means that for both `torch.add(a, b)`, `torch.stack([a, b])` (and
+        Recursively flatten input and yield all instances of `MetaObj`.
+        This means that for both `torch.add(a, b)`, `torch.stack([a, b])` (and
         their numpy equivalents), we return `[a, b]` if both `a` and `b` are of type
         `MetaObj`.
 
         Args:
-            args: Sequence of inputs to be flattened.
+            args: Iterables of inputs to be flattened.
         Returns:
             list of nested `MetaObj` from input.
         """
-        out = []
-        for a in args:
+        for a in itertools.chain(*args):
             if isinstance(a, (list, tuple)):
-                out += MetaObj.flatten_meta_objs(a)
+                yield from MetaObj.flatten_meta_objs(a)
             elif isinstance(a, MetaObj):
-                out.append(a)
-        return out
+                yield a
 
-    def _copy_attr(self, attribute: str, input_objs: list[MetaObj], default_fn: Callable, deep_copy: bool) -> None:
+    def _copy_attr(self, attributes: list[str], input_objs, defaults: list, deep_copy: bool) -> None:
         """
-        Copy an attribute from the first in a list of `MetaObj`. In the case of
+        Copy attributes from the first in a list of `MetaObj`. In the case of
         `torch.add(a, b)`, both `a` and `b` could be `MetaObj` or something else, so
         check them all. Copy the first to `self`.
 
         We also perform a deep copy of the data if desired.
 
         Args:
-            attribute: string corresponding to attribute to be copied (e.g., `meta`).
-            input_objs: List of `MetaObj`. We'll copy the attribute from the first one
+            attributes: a sequence of strings corresponding to attributes to be copied (e.g., `['meta']`).
+            input_objs: an iterable of `MetaObj` instances. We'll copy the attribute from the first one
                 that contains that particular attribute.
-            default_fn: If none of `input_objs` have the attribute that we're
-                interested in, then use this default function (e.g., `lambda: {}`.)
-            deep_copy: Should the attribute be deep copied? See `_copy_meta`.
+            defaults: If none of `input_objs` have the attribute that we're
+                interested in, then use this default value/function (e.g., `lambda: {}`.)
+                the defaults must be the same length as `attributes`.
+            deep_copy: whether to deep copy the corresponding attribute.
 
         Returns:
             Returns `None`, but `self` should be updated to have the copied attribute.
         """
-        attributes = [getattr(i, attribute) for i in input_objs if hasattr(i, attribute)]
-        if len(attributes) > 0:
-            val = attributes[0]
-            if deep_copy:
-                val = deepcopy(val)
-            setattr(self, attribute, val)
-        else:
-            setattr(self, attribute, default_fn())
-
-    def _copy_meta(self, input_objs: list[MetaObj]) -> None:
+        found = [False] * len(attributes)
+        for i, (idx, a) in itertools.product(input_objs, enumerate(attributes)):
+            if not found[idx] and hasattr(i, a):
+                setattr(self, a, deepcopy(getattr(i, a)) if deep_copy else getattr(i, a))
+                found[idx] = True
+            if all(found):
+                return
+        for a, f, d in zip(attributes, found, defaults):
+            if not f:
+                setattr(self, a, d() if callable(defaults) else d)
+        return
+
+    def _copy_meta(self, input_objs, deep_copy=False) -> None:
         """
-        Copy metadata from a list of `MetaObj`. For a given attribute, we copy the
+        Copy metadata from an iterable of `MetaObj` instances. For a given attribute, we copy the
         adjunct data from the first element in the list containing that attribute.
 
-        If there has been a change in `id` (e.g., `a=b+c`), then deepcopy. Else (e.g.,
-        `a+=1`), then don't.
-
         Args:
             input_objs: list of `MetaObj` to copy data from.
 
         """
-        id_in = id(input_objs[0]) if len(input_objs) > 0 else None
-        deep_copy = id(self) != id_in
-        self._copy_attr("meta", input_objs, self.get_default_meta, deep_copy)
-        self._copy_attr("applied_operations", input_objs, self.get_default_applied_operations, deep_copy)
-        self.is_batch = input_objs[0].is_batch if len(input_objs) > 0 else False
+        self._copy_attr(
+            ["meta", "applied_operations"],
+            input_objs,
+            [MetaObj.get_default_meta(), MetaObj.get_default_applied_operations()],
+            deep_copy,
+        )
 
-    def get_default_meta(self) -> dict:
+    @staticmethod
+    def get_default_meta() -> dict:
         """Get the default meta.
 
         Returns:
             default metadata.
         """
         return {}
 
-    def get_default_applied_operations(self) -> list:
+    @staticmethod
+    def get_default_applied_operations() -> list:
         """Get the default applied operations.
 
         Returns:
@@ -183,28 +186,28 @@ def __repr__(self) -> str:
     @property
     def meta(self) -> dict:
         """Get the meta."""
-        return self._meta if hasattr(self, "_meta") else self.get_default_meta()
+        return self._meta if hasattr(self, "_meta") else MetaObj.get_default_meta()
 
     @meta.setter
     def meta(self, d) -> None:
         """Set the meta."""
         if d == TraceKeys.NONE:
-            self._meta = self.get_default_meta()
+            self._meta = MetaObj.get_default_meta()
         self._meta = d
 
     @property
     def applied_operations(self) -> list:
         """Get the applied operations."""
         if hasattr(self, "_applied_operations"):
             return self._applied_operations
-        return self.get_default_applied_operations()
+        return MetaObj.get_default_applied_operations()
 
     @applied_operations.setter
     def applied_operations(self, t) -> None:
         """Set the applied operations."""
         if t == TraceKeys.NONE:
             # received no operations when decollating a batch
-            self._applied_operations = self.get_default_applied_operations()
+            self._applied_operations = MetaObj.get_default_applied_operations()
             return
         self._applied_operations = t
 

monai/data/meta_tensor.py

@@ -13,7 +13,7 @@
 
 import warnings
 from copy import deepcopy
-from typing import Any, Callable, Sequence
+from typing import Any, Sequence
 
 import torch
 
@@ -126,19 +126,20 @@ def __init__(
         elif isinstance(x, MetaTensor):
             self.applied_operations = x.applied_operations
         else:
-            self.applied_operations = self.get_default_applied_operations()
+            self.applied_operations = MetaObj.get_default_applied_operations()
 
         # if we are creating a new MetaTensor, then deep copy attributes
         if isinstance(x, torch.Tensor) and not isinstance(x, MetaTensor):
             self.meta = deepcopy(self.meta)
             self.applied_operations = deepcopy(self.applied_operations)
         self.affine = self.affine.to(self.device)
 
-    def _copy_attr(self, attribute: str, input_objs: list[MetaObj], default_fn: Callable, deep_copy: bool) -> None:
-        super()._copy_attr(attribute, input_objs, default_fn, deep_copy)
-        val = getattr(self, attribute)
-        if isinstance(val, torch.Tensor):
-            setattr(self, attribute, val.to(self.device))
+    def _copy_attr(self, attributes: list[str], input_objs, defaults: list, deep_copy: bool) -> None:
+        super()._copy_attr(attributes, input_objs, defaults, deep_copy)
+        for a in attributes:
+            val = getattr(self, a)
+            if isinstance(val, torch.Tensor):
+                setattr(self, a, val.to(self.device))
 
     @staticmethod
     def update_meta(rets: Sequence, func, args, kwargs) -> Sequence:
@@ -173,6 +174,7 @@ def update_meta(rets: Sequence, func, args, kwargs) -> Sequence:
         """
         out = []
         metas = None
+        is_batch = any(x.is_batch for x in MetaObj.flatten_meta_objs(args, kwargs.values()) if hasattr(x, "is_batch"))
         for idx, ret in enumerate(rets):
             # if not `MetaTensor`, nothing to do.
             if not isinstance(ret, MetaTensor):
@@ -182,20 +184,18 @@ def update_meta(rets: Sequence, func, args, kwargs) -> Sequence:
                 ret = ret.as_tensor()
             # else, handle the `MetaTensor` metadata.
             else:
-                meta_args = MetaObj.flatten_meta_objs(list(args) + list(kwargs.values()))
-                # this is not implemented but the network arch may run into this case:
+                meta_args = MetaObj.flatten_meta_objs(args, kwargs.values())  # type: ignore
+                ret._copy_meta(meta_args, deep_copy=not is_batch)
+                ret.is_batch = is_batch
+                # the following is not implemented but the network arch may run into this case:
                 # if func == torch.cat and any(m.is_batch if hasattr(m, "is_batch") else False for m in meta_args):
                 #     raise NotImplementedError("torch.cat is not implemented for batch of MetaTensors.")
-                ret._copy_meta(meta_args)
 
                 # If we have a batch of data, then we need to be careful if a slice of
                 # the data is returned. Depending on how the data are indexed, we return
                 # some or all of the metadata, and the return object may or may not be a
                 # batch of data (e.g., `batch[:,-1]` versus `batch[0]`).
-                if ret.is_batch:
-                    # only decollate metadata once
-                    if metas is None:
-                        metas = decollate_batch(ret.meta)
+                if is_batch:
                     # if indexing e.g., `batch[0]`
                     if func == torch.Tensor.__getitem__:
                         batch_idx = args[1]
@@ -205,6 +205,9 @@ def update_meta(rets: Sequence, func, args, kwargs) -> Sequence:
                         # first element will be `slice(None, None, None)` and `Ellipsis`,
                         # respectively. Don't need to do anything with the metadata.
                         if batch_idx not in (slice(None, None, None), Ellipsis):
+                            # only decollate metadata once
+                            if metas is None:
+                                metas = decollate_batch(ret.meta)
                             meta = metas[batch_idx]
                             # if using e.g., `batch[0:2]`, then `is_batch` should still be
                             # `True`. Also re-collate the remaining elements.
@@ -226,6 +229,8 @@ def update_meta(rets: Sequence, func, args, kwargs) -> Sequence:
                         else:
                             dim = 0
                         if dim == 0:
+                            if metas is None:
+                                metas = decollate_batch(ret.meta)
                             ret.meta = metas[idx]
                             ret.is_batch = False
 

tests/test_integration_fast_train.py

@@ -34,7 +34,6 @@
     Compose,
     CropForegroundd,
     EnsureChannelFirstd,
-    EnsureTyped,
     FgBgToIndicesd,
     LoadImaged,
     RandAffined,
@@ -94,7 +93,6 @@ def test_train_timing(self):
                 # and cache them to accelerate training
                 FgBgToIndicesd(keys="label", fg_postfix="_fg", bg_postfix="_bg"),
                 # move the data to GPU and cache to avoid CPU -> GPU sync in every epoch
-                EnsureTyped(keys=["image", "label"], drop_meta=True),
                 ToDeviced(keys=["image", "label"], device=device),
                 # randomly crop out patch samples from big
                 # image based on pos / neg ratio