update apply_to_collections to support dataclass inputs (#11889)

armanal · rohitgr7 · web-flow · commit 834115b85910 · 2022-03-28T18:44:13.000+02:00
Co-authored-by: rohitgr7 &lt;rohitgr1998@gmail.com&gt;
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -172,6 +172,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 - Added support for Habana Accelerator (HPU) ([#11808](https://github.com/PyTorchLightning/pytorch-lightning/pull/11808))
 
+- Added support for dataclasses in `apply_to_collections` ([#11889](https://github.com/PyTorchLightning/pytorch-lightning/pull/11889))
+
+
 
 ### Changed
 
diff --git a/pytorch_lightning/utilities/apply_func.py b/pytorch_lightning/utilities/apply_func.py
@@ -138,6 +138,7 @@ def apply_to_collection(
         result = deepcopy(data, memo=memo)
         # apply function to each field
         for field_name, (field_value, field_init) in fields.items():
+            v = None
             if field_init:
                 v = apply_to_collection(
                     field_value,
@@ -215,13 +216,65 @@ def apply_to_collections(
     is_namedtuple = _is_namedtuple(data1)
     is_sequence = isinstance(data1, Sequence) and not isinstance(data1, str)
     if (is_namedtuple or is_sequence) and data2 is not None:
-        assert len(data1) == len(data2), "Sequence collections have different sizes"
+        assert len(data1) == len(data2), "Sequence collections have different sizes."
         out = [
             apply_to_collections(v1, v2, dtype, function, *args, wrong_dtype=wrong_dtype, **kwargs)
             for v1, v2 in zip(data1, data2)
         ]
         return elem_type(*out) if is_namedtuple else elem_type(out)
 
+    if _is_dataclass_instance(data1) and data2 is not None:
+        if not _is_dataclass_instance(data2):
+            raise TypeError(
+                "Expected inputs to be dataclasses of the same type or to have identical fields"
+                f" but got input 1 of type {type(data1)} and input 2 of type {type(data2)}."
+            )
+        if not (
+            len(dataclasses.fields(data1)) == len(dataclasses.fields(data2))
+            and all(map(lambda f1, f2: isinstance(f1, type(f2)), dataclasses.fields(data1), dataclasses.fields(data2)))
+        ):
+            raise TypeError("Dataclasses fields do not match.")
+        # make a deepcopy of the data,
+        # but do not deepcopy mapped fields since the computation would
+        # be wasted on values that likely get immediately overwritten
+        data = [data1, data2]
+        fields: List[dict] = [{}, {}]
+        memo: dict = {}
+        for i in range(len(data)):
+            for field in dataclasses.fields(data[i]):
+                field_value = getattr(data[i], field.name)
+                fields[i][field.name] = (field_value, field.init)
+                if i == 0:
+                    memo[id(field_value)] = field_value
+
+        result = deepcopy(data1, memo=memo)
+
+        # apply function to each field
+        for ((field_name, (field_value1, field_init1)), (_, (field_value2, field_init2))) in zip(
+            fields[0].items(), fields[1].items()
+        ):
+            v = None
+            if field_init1 and field_init2:
+                v = apply_to_collections(
+                    field_value1,
+                    field_value2,
+                    dtype,
+                    function,
+                    *args,
+                    wrong_dtype=wrong_dtype,
+                    **kwargs,
+                )
+            if not field_init1 or not field_init2 or v is None:  # retain old value
+                return apply_to_collection(data1, dtype, function, *args, wrong_dtype=wrong_dtype, **kwargs)
+            try:
+                setattr(result, field_name, v)
+            except dataclasses.FrozenInstanceError as e:
+                raise MisconfigurationException(
+                    "A frozen dataclass was passed to `apply_to_collections` but this is not allowed."
+                    " HINT: is your batch a frozen dataclass?"
+                ) from e
+        return result
+
     return apply_to_collection(data1, dtype, function, *args, wrong_dtype=wrong_dtype, **kwargs)
 
 
diff --git a/tests/utilities/test_apply_func.py b/tests/utilities/test_apply_func.py
@@ -25,89 +25,94 @@
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 
 
+@dataclasses.dataclass
+class Feature:
+    input_ids: torch.Tensor
+    segment_ids: np.ndarray
+
+    def __eq__(self, o: object) -> bool:
+        if not isinstance(o, Feature):
+            return NotImplemented
+
+        return torch.equal(self.input_ids, o.input_ids) and np.equal(self.segment_ids, o.segment_ids).all()
+
+
+@dataclasses.dataclass
+class ModelExample:
+    example_ids: List[str]
+    feature: Feature
+    label: torch.Tensor
+    some_constant: int = dataclasses.field(init=False)
+
+    def __post_init__(self):
+        self.some_constant = 7
+
+    def __eq__(self, o: object) -> bool:
+        if not isinstance(o, ModelExample):
+            return NotImplemented
+
+        return (
+            self.example_ids == o.example_ids
+            and self.feature == o.feature
+            and torch.equal(self.label, o.label)
+            and self.some_constant == o.some_constant
+        )
+
+
+@dataclasses.dataclass
+class WithClassVar:
+    class_var: ClassVar[int] = 0
+    dummy: Any
+
+    def __eq__(self, o: object) -> bool:
+        if not isinstance(o, WithClassVar):
+            return NotImplemented
+        elif isinstance(self.dummy, torch.Tensor):
+            return torch.equal(self.dummy, o.dummy)
+
+        return self.dummy == o.dummy
+
+
+@dataclasses.dataclass
+class WithInitVar:
+    dummy: Any
+    override: InitVar[Optional[Any]] = None
+
+    def __post_init__(self, override: Optional[Any]):
+        if override is not None:
+            self.dummy = override
+
+    def __eq__(self, o: object) -> bool:
+        if not isinstance(o, WithInitVar):
+            return NotImplemented
+        elif isinstance(self.dummy, torch.Tensor):
+            return torch.equal(self.dummy, o.dummy)
+
+        return self.dummy == o.dummy
+
+
+@dataclasses.dataclass
+class WithClassAndInitVar:
+    class_var: ClassVar[torch.Tensor] = torch.tensor(0)
+    dummy: Any
+    override: InitVar[Optional[Any]] = torch.tensor(1)
+
+    def __post_init__(self, override: Optional[Any]):
+        if override is not None:
+            self.dummy = override
+
+    def __eq__(self, o: object) -> bool:
+        if not isinstance(o, WithClassAndInitVar):
+            return NotImplemented
+        elif isinstance(self.dummy, torch.Tensor):
+            return torch.equal(self.dummy, o.dummy)
+
+        return self.dummy == o.dummy
+
+
 def test_recursive_application_to_collection():
     ntc = namedtuple("Foo", ["bar"])
 
-    @dataclasses.dataclass
-    class Feature:
-        input_ids: torch.Tensor
-        segment_ids: np.ndarray
-
-        def __eq__(self, o: object) -> bool:
-            if not isinstance(o, Feature):
-                return NotImplemented
-            else:
-                return torch.equal(self.input_ids, o.input_ids) and np.equal(self.segment_ids, o.segment_ids).all()
-
-    @dataclasses.dataclass
-    class ModelExample:
-        example_ids: List[str]
-        feature: Feature
-        label: torch.Tensor
-        some_constant: int = dataclasses.field(init=False)
-
-        def __post_init__(self):
-            self.some_constant = 7
-
-        def __eq__(self, o: object) -> bool:
-            if not isinstance(o, ModelExample):
-                return NotImplemented
-            else:
-                return (
-                    self.example_ids == o.example_ids
-                    and self.feature == o.feature
-                    and torch.equal(self.label, o.label)
-                    and self.some_constant == o.some_constant
-                )
-
-    @dataclasses.dataclass
-    class WithClassVar:
-        class_var: ClassVar[int] = 0
-        dummy: Any
-
-        def __eq__(self, o: object) -> bool:
-            if not isinstance(o, WithClassVar):
-                return NotImplemented
-            elif isinstance(self.dummy, torch.Tensor):
-                return torch.equal(self.dummy, o.dummy)
-            else:
-                return self.dummy == o.dummy
-
-    @dataclasses.dataclass
-    class WithInitVar:
-        dummy: Any
-        override: InitVar[Optional[Any]] = None
-
-        def __post_init__(self, override: Optional[Any]):
-            if override is not None:
-                self.dummy = override
-
-        def __eq__(self, o: object) -> bool:
-            if not isinstance(o, WithInitVar):
-                return NotImplemented
-            elif isinstance(self.dummy, torch.Tensor):
-                return torch.equal(self.dummy, o.dummy)
-            else:
-                return self.dummy == o.dummy
-
-    @dataclasses.dataclass
-    class WithClassAndInitVar:
-        class_var: ClassVar[torch.Tensor] = torch.tensor(0)
-        dummy: Any
-        override: InitVar[Optional[Any]] = torch.tensor(1)
-
-        def __post_init__(self, override: Optional[Any]):
-            if override is not None:
-                self.dummy = override
-
-        def __eq__(self, o: object) -> bool:
-            if not isinstance(o, WithClassAndInitVar):
-                return NotImplemented
-            elif isinstance(self.dummy, torch.Tensor):
-                return torch.equal(self.dummy, o.dummy)
-            else:
-                return self.dummy == o.dummy
-
     model_example = ModelExample(
         example_ids=["i-1", "i-2", "i-3"],
         feature=Feature(input_ids=torch.tensor([1.0, 2.0, 3.0]), segment_ids=np.array([4.0, 5.0, 6.0])),
@@ -303,6 +308,36 @@ def fn(a, b):
     assert reduced is None
 
 
+def test_apply_to_collections_dataclass():
+    to_reduce_1 = Feature(input_ids=torch.tensor([1.0, 2.0, 3.0]), segment_ids=np.array([4.0, 5.0, 6.0]))
+    to_reduce_2 = Feature(input_ids=torch.tensor([1.0, 2.0, 3.0]), segment_ids=np.array([4.0, 5.0, 6.0]))
+
+    def fn(a, b):
+        return a + b
+
+    reduced = apply_to_collections(to_reduce_1, to_reduce_2, (torch.Tensor, numbers.Number, np.ndarray), fn)
+
+    assert reduced == Feature(input_ids=torch.tensor([2.0, 4.0, 6.0]), segment_ids=np.array([8.0, 10.0, 12.0]))
+
+    model_example = ModelExample(
+        example_ids=["i-1", "i-2", "i-3"],
+        feature=to_reduce_1,
+        label=torch.tensor([7.0, 8.0, 9.0]),
+    )
+
+    # different types
+    with pytest.raises(TypeError, match="Expected inputs to be dataclasses of the same type"):
+        apply_to_collections(to_reduce_1, [1, 2], (torch.Tensor, numbers.Number, np.ndarray), fn)
+
+    # unmatched fields
+    with pytest.raises(TypeError, match="Dataclasses fields do not match"):
+        apply_to_collections(to_reduce_1, model_example, (torch.Tensor, numbers.Number, np.ndarray), fn)
+
+    classvar = WithClassVar(torch.arange(3))  # dataclass with same number but different type of fields
+    with pytest.raises(TypeError, match="Dataclasses fields do not match"):
+        apply_to_collections(to_reduce_1, classvar, (torch.Tensor, numbers.Number, np.ndarray), fn)
+
+
 def test_apply_to_collection_frozen_dataclass():
     @dataclasses.dataclass(frozen=True)
     class Foo:
