Repository: Lightning-AI/pytorch-lightning

Issue:
Breaking compatibility with custom datatypes implementing `.to`
## 🚀 Feature
<!-- A clear and concise description of the feature proposal -->
Bring back compatibility for custom datatypes in collections implementing `.to` for transferring data.

### Motivation

<!-- Please outline the motivation for the proposal. Is your feature request related to a problem? e.g., I'm always frustrated when [...]. If this is related to another GitHub issue, please link here too -->
I am using Pytorch Lightning together with Pytorch Geometric. Pytorch Geometric implements several custom datatypes and dataloaders which is really useful for geometric deep learning. Everything worked well with pytorch lightning 0.7.6, as the custom datatypes implement a `.to` method for transferring the data to different devices.
However, with the recent 0.8.1 update, this is no longer possible and I had to scour the documentation to be able to implement a fix using `transfer_batch_to_device(batch, device)`. This is in my opinion not very pretty, as my batch looks like this 
```
{"data": pytorch geometric batch object, "id": tensor, ...}
```
i.e. it is just a dictionary of types that all implement the `.to` method.

### Pitch

<!-- A clear and concise description of what you want to happen. -->

- Make it possible for classes implementing the `.to` method to be transferred automatically
- If part of the batch could not be transferred automatically output a warning letting the user know, that a custom transfer function for the batch might be required, or to implement the `.to` method for custom datatypes in the batch
- Add a note to the introduction guide about custom datatypes and handling for custom datatypes

### Alternatives

<!-- A clear and concise description of any alternative solutions or features you've considered, if any. -->
If this change was intentional and the behavior of trying to call the `.to` method is not desired, I think there should definitely be some more documentation about this, in a more obvious place.

### Additional context

<!-- Add any other context or screenshots about the feature request here. -->



Discussion:
Hi! thanks for your contribution!, great first issue!
This should definitely work. It was not my intention to break it and I thought I had even a test for it. I will look into this and also converting the issue to a bug report.
Have identified the problem and fix is in the works... Silly oversight on my part.

Regarding the warning, not sure. Is it not enough to point out in the docs that anything not a tensor (or .to() movable) will be left untouched unless the hook is implemented? 

Otherwise, such a check would have to be done on every batch, because in theory the batch can (structurally) change every iteration, and that would cause a spam of warning messages. 
> Have identified the problem and fix is in the works... Silly oversight on my part.

Awesome, thank you!

> Regarding the warning, not sure. Is it not enough to point out in the docs that anything not a tensor (or .to() movable) will be left untouched unless the hook is implemented?
>
> Otherwise, such a check would have to be done on every batch, because in theory the batch can (structurally) change every iteration, and that would cause a spam of warning messages.

That's a good point, in that case I think the docs should suffice. Though maybe the compatible data types can be mentioned in the Introduction Guide? Reading through the Introduction Guide I found no mention of this anywhere, and I think it may fit well into the Data Section. Possibly after the *Models defined by data* section, there could be a section which talks briefly about custom data loaders, custom batches and what kind of custom data lightning can handle automatically. Maybe also just refer to the appropriate place in the documentation.

Resolving patch:
diff --git a/pytorch_lightning/core/hooks.py b/pytorch_lightning/core/hooks.py
--- a/pytorch_lightning/core/hooks.py
+++ b/pytorch_lightning/core/hooks.py
@@ -204,7 +204,7 @@ def transfer_batch_to_device(self, batch: Any, device: torch.device) -> Any:
 
         The data types listed below (and any arbitrary nesting of them) are supported out of the box:
 
-        - :class:`torch.Tensor`
+        - :class:`torch.Tensor` or anything that implements `.to(...)`
         - :class:`list`
         - :class:`dict`
         - :class:`tuple`
diff --git a/pytorch_lightning/utilities/apply_func.py b/pytorch_lightning/utilities/apply_func.py
--- a/pytorch_lightning/utilities/apply_func.py
+++ b/pytorch_lightning/utilities/apply_func.py
@@ -1,3 +1,4 @@
+from abc import ABC
 from collections import Mapping, Sequence
 from typing import Any, Callable, Union
 
@@ -38,14 +39,43 @@ def apply_to_collection(data: Any, dtype: Union[type, tuple], function: Callable
     return data
 
 
+class TransferableDataType(ABC):
+    """
+    A custom type for data that can be moved to a torch device via `.to(...)`.
+
+    Example:
+
+        >>> isinstance(dict, TransferableDataType)
+        False
+        >>> isinstance(torch.rand(2, 3), TransferableDataType)
+        True
+        >>> class CustomObject:
+        ...     def __init__(self):
+        ...         self.x = torch.rand(2, 2)
+        ...     def to(self, device):
+        ...         self.x = self.x.to(device)
+        ...         return self
+        >>> isinstance(CustomObject(), TransferableDataType)
+        True
+    """
+
+    @classmethod
+    def __subclasshook__(cls, subclass):
+        if cls is TransferableDataType:
+            to = getattr(subclass, "to", None)
+            return callable(to)
+        return NotImplemented
+
+
 def move_data_to_device(batch: Any, device: torch.device):
     """
-    Transfers a collection of tensors to the given device.
+    Transfers a collection of data to the given device. Any object that defines a method
+    ``to(device)`` will be moved and all other objects in the collection will be left untouched.
 
     Args:
-        batch: A tensor or collection of tensors. See :func:`apply_to_collection`
-            for a list of supported collection types.
-        device: The device to which tensors should be moved
+        batch: A tensor or collection of tensors or anything that has a method `.to(...)`.
+            See :func:`apply_to_collection` for a list of supported collection types.
+        device: The device to which the data should be moved
 
     Return:
         the same collection but with all contained tensors residing on the new device.
@@ -54,6 +84,6 @@ def move_data_to_device(batch: Any, device: torch.device):
         - :meth:`torch.Tensor.to`
         - :class:`torch.device`
     """
-    def to(tensor):
-        return tensor.to(device, non_blocking=True)
-    return apply_to_collection(batch, dtype=torch.Tensor, function=to)
+    def to(data):
+        return data.to(device, non_blocking=True)
+    return apply_to_collection(batch, dtype=TransferableDataType, function=to)


Test patch:
diff --git a/tests/models/test_gpu.py b/tests/models/test_gpu.py
--- a/tests/models/test_gpu.py
+++ b/tests/models/test_gpu.py
@@ -286,3 +286,15 @@ def test_single_gpu_batch_parse():
     batch = trainer.transfer_batch_to_gpu(batch, 0)
     assert batch[0].a.device.index == 0
     assert batch[0].a.type() == 'torch.cuda.FloatTensor'
+
+    # non-Tensor that has `.to()` defined
+    class CustomBatchType:
+        def __init__(self):
+            self.a = torch.rand(2, 2)
+
+        def to(self, *args, **kwargs):
+            self.a = self.a.to(*args, **kwargs)
+            return self
+
+    batch = trainer.transfer_batch_to_gpu(CustomBatchType())
+    assert batch.a.type() == 'torch.cuda.FloatTensor'