Repository: Lightning-AI/pytorch-lightning

Issue:
Improve moving data / model to GPU using torchtext
## 🚀 Feature
Improve moving data and model to GPU if torchtext is used.

### Motivation
#### Case 1:
`Batch` object generated by `torchtext.data.Iterator` doesn't follow the rules described here https://github.com/PyTorchLightning/pytorch-lightning/blob/45d671a4a81788b9d97fd6b47763816926e58e95/pytorch_lightning/trainer/distrib_parts.py#L420

As the result data is not moved to GPU. `torchtext.data.Iterator` is returned by method `train_dataloader`. Take in mind that `torchtext.data.Iterator` has a device argument that is not properly utilized by pytorch-ligthning.

```
    @ptl.data_loader
    def train_dataloader(self):
        ...  
        return Iterator(dataset=dataset, batch_size=self.batch_size, shuffle=False, device=DEVICE)
```

Partially reported here https://github.com/PyTorchLightning/pytorch-lightning/issues/226

#### Case 2
Using torchtext you can read pre-trained embeddings and create `nn.Embedding` object as follows

```
    def train_dataloader(self):
        ...
        self.text_field.build_vocab(
            dataset,
            vectors=Vectors("/data/embeddings/glove/glove.840B.300d.txt"),
        )

        self.embeddings = nn.Embedding(
            ...
            padding_idx=self.text_field.vocab.stoi[PAD_TOKEN],
            _weight=self.text_field.vocab.vectors.to(DEVICE),
        )
```

`nn.Embedding` is clearly dependent on `self.text_field.vocab` and this is in turn dependent on `dataset` that is used by `train_dataloader`. Currently any part of the model that is not created fully in `__init__` of the `ptl.LigthningModule` is not moved to the GPU. It requires still to have a global variable that determines a device i.e. `DEVICE`. It makes `Trainer(n_gpus=...)` useless.

### Pitch
I would like not to worry about moving data to GPU using torchtext combined with pytorch-lightning.


Discussion:
@mateuszpieniak which version of lightning are you actually using?
@PyTorchLightning/core-contributors suggestions? ^^
@Borda Version == 0.6.0, but I see that there is some new release ^^ Let me check it for the new version.

* Some solution for Case 1 could be not only looking for `.to()` or `.cuda()`, attribute, but also `.device` field inside the class? 
* Some solution for Case 2 could be moving the model to GPU a bit further in "code but it is not so straightforward since 

https://github.com/PyTorchLightning/pytorch-lightning/blob/45d671a4a81788b9d97fd6b47763816926e58e95/pytorch_lightning/trainer/distrib_parts.py#L457

moving the model to GPU and configures AMP straightaway. 

@Borda Well, for version == 0.7.1 I have a slightly different issue.

As I mentioned before, I build a vocabulary using the training dataset in `train_dataloader`. 
```
    def train_dataloader(self):
        ...
        self.text_field.build_vocab(
            dataset,
            vectors=Vectors("/data/embeddings/glove/glove.840B.300d.txt"),
        )
```
It creates a `self.text_field.vocab` object field. I use the same `self.text_field` (class `torchtext.data.Field`) for `val_dataloader`. It is expected since I want to have the same preprocessing / tokenization / embeddings for validation and test set. This time I got

```
AttributeError: 'Field' object has no attribute 'vocab'
```

because of running sanity check for the validation set before training happens (or to be precise before `train_dataloader` happens) - https://github.com/PyTorchLightning/pytorch-lightning/blob/master/pytorch_lightning/trainer/trainer.py#L883

I thought that I could disable this using `fast_dev_run`, but it seems like such condition doesn't work properly - https://github.com/PyTorchLightning/pytorch-lightning/blob/master/pytorch_lightning/trainer/trainer.py#L879
@neggert @srush pls ^^
@Borda I am still not able to get torchtext iterators to work with pytorch lightning even on a single GPU as @mateuszpieniak mentioned in his first Case above and in #226 

The (very bad) workaround which I have been using so far is to manually move the data to a GPU while creating the iterator (by passing the device to the torchtext iterator). This prevents me from utilizing multiple GPUs and also unnecessarily loads all the data on the GPU in the beginning. Has there been any progress on this issue? 

Thanks 
Sorry, don't know torchtext well, but if I understand correctly, the problem is that the internal method that moves the data to gpu checks for some known types (tuple, list, dict) but when passing a custom batch object, it does not know how to do that. 

My pitch: 
Introduce a model hook (or callback, not sure?). 
`__transfer_data_to_device ` will call the hook first:

```python
def __transfer_data_to_device(...)
    
    batch = model.transfer_batch_to_device(batch, device)  # call the hook
    if batch is not None:  # for example, it returns None if user did not override it
        return batch

    # otherwise continue with built in detection of tuples etc.
```
This allows anyone to implement their own transfers for custom objects. 
Maybe we should still support the torchtext out of the box, but if we can't, this could be an option. 
@mateuszpieniak would such a model hook help you solve the problem with the dataloading and moving?
@awaelchli You understood the case 1 correctly. Well, it should work :smiley: I wasn't aware of the hook concept in PyTorchLightning. I'll let you know whether it works. Thanks!
Case 1) @awaelchli I ended up with creating my custom Trainer and overriding a private method:
```
    def _TrainerDPMixin__transfer_data_to_device(self, batch, device, gpu_id=None):
        """ Handles torchtext.data.Batch. """
        ...
        # when torchtext.data.Batch
        if isinstance(batch, Batch):
            for field in batch.input_fields + batch.target_fields:
                data = self._TrainerDPMixin__transfer_data_to_device(getattr(batch, field), device, gpu_id)
                setattr(batch, field, data)
            return batch
```
Case 2) I just build dictionary in the `__init__` method.

If we don't plan to support torchtext, then then I think we can consider this issue closed.
Case 1: that certainly works. If we introduced a hook for that, then the user of PL would not need to modify the source code of Trainer. It is not meant to be modified. 

Case 2: If you want a tensor in your model be moved with the model (when .to(..) is called on the model) then register it as a buffer:
https://pytorch.org/docs/stable/nn.html?highlight=register_buffer#torch.nn.Module.register_buffer
It can be accessed just like a regular attribute of the model, but there is no need to pass in the device. 
> If we don't plan to support torchtext, then then I think we can consider this issue closed.

I would not close it yet. @PyTorchLightning/core-contributors should torchtext be supported directly in PL or do we provide a hook as I suggested, or both?
@williamFalcon ^^
1. let’s introduce s formal hook
2. let’s formally support torchtext! 

good ideas everyone! 
@awaelchli i would prefer a hook since that will be generally useful and will likely reduce the long-term maintenance burden
@mateuszpieniak the model hook I added should make it easier to work with the Batch object in torchtext, but I have not actually tested it yet in a full example.
Do you by chance have already have mini example for torchtext with PL? I have 0 experience with torchtext, but will look into it to provide better support.
@awaelchli Super simple classification example with a bag of embeddings. It should work both on GPU & CPU.

https://gist.github.com/mateuszpieniak/f290b3a727db7e94b9da0bd3bd2e33c1
great! i can start with that!
nice, @mateuszpieniak or if you want to take is and @awaelchli may assist you :]
@Borda Sure, I will do this over the weekend unless @awaelchli already started to work on this :]
Not yet started, and I don't mind you taking it if you are interested.
Great discussion. 
Currently, is it possible or not to use PyTorch Lightning in multi-GPUs training?
? of course it's possible... this is at the core of lightning...

If you specifically mean for torchtext, then it sounds like there's some edge case it introduces that we are working on.
@mateuszpieniak we're preparing 0.8.2 now. Want to submit this PR?
> ? of course it's possible... this is at the core of lightning...
> 
> If you specifically mean for torchtext, then it sounds like there's some edge case it introduces that we are working on.

Sorry, indeed I meant multi-GPU training with torchtext.

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
@@ -208,7 +208,7 @@ def transfer_batch_to_device(self, batch: Any, device: torch.device) -> Any:
         - :class:`list`
         - :class:`dict`
         - :class:`tuple`
-        - ``torchtext.data.Batch`` (COMING SOON)
+        - :class:`torchtext.data.batch.Batch`
 
         For anything else, you need to define how the data is moved to the target device (CPU, GPU, TPU, ...).
 
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
@@ -3,6 +3,8 @@
 from typing import Any, Callable, Union
 
 import torch
+from torchtext.data import Batch
+from copy import copy
 
 
 def apply_to_collection(data: Any, dtype: Union[type, tuple], function: Callable, *args, **kwargs) -> Any:
@@ -84,6 +86,16 @@ def move_data_to_device(batch: Any, device: torch.device):
         - :meth:`torch.Tensor.to`
         - :class:`torch.device`
     """
-    def to(data):
+
+    def batch_to(data):
+        if isinstance(data, Batch):
+            # Shallow copy because each Batch has a reference to Dataset which contains all examples
+            device_data = copy(data)
+            for field in data.fields:
+                # Batch contains output of Field.process(...) which is tensor hence .to(...) exists
+                device_field = getattr(data, field).to(device, non_blocking=True)
+                setattr(device_data, field, device_field)
+            return device_data
+
         return data.to(device, non_blocking=True)
-    return apply_to_collection(batch, dtype=TransferableDataType, function=to)
+    return apply_to_collection(batch, dtype=(TransferableDataType, Batch), function=batch_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
@@ -9,6 +9,7 @@
 from pytorch_lightning.trainer.distrib_parts import _parse_gpu_ids, determine_root_gpu_device
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from tests.base import EvalModelTemplate
+from torchtext.data import Batch, Dataset, Example, Field, LabelField
 
 PRETEND_N_OF_GPUS = 16
 
@@ -298,3 +299,32 @@ def to(self, *args, **kwargs):
 
     batch = trainer.transfer_batch_to_gpu(CustomBatchType())
     assert batch.a.type() == 'torch.cuda.FloatTensor'
+
+    # torchtext.data.Batch
+    samples = [
+        {'text': 'PyTorch Lightning is awesome!', 'label': 0},
+        {'text': 'Please make it work with torchtext', 'label': 1}
+    ]
+
+    text_field = Field()
+    label_field = LabelField()
+    fields = {
+        'text': ('text', text_field),
+        'label': ('label', label_field)
+    }
+
+    examples = [Example.fromdict(sample, fields) for sample in samples]
+    dataset = Dataset(
+        examples=examples,
+        fields=fields.values()
+    )
+
+    # Batch runs field.process() that numericalizes tokens, but it requires to build dictionary first
+    text_field.build_vocab(dataset)
+    label_field.build_vocab(dataset)
+
+    batch = Batch(data=examples, dataset=dataset)
+    batch = trainer.transfer_batch_to_gpu(batch, 0)
+
+    assert batch.text.type() == 'torch.cuda.LongTensor'
+    assert batch.label.type() == 'torch.cuda.LongTensor'