Repository: Lightning-AI/pytorch-lightning

Issue:
to() got an unexpected keyword argument 'non_blocking' for DGLGraph
<!-- 
### Common bugs:
1. Tensorboard not showing in Jupyter-notebook see [issue 79](https://github.com/PyTorchLightning/pytorch-lightning/issues/79).    
2. PyTorch 1.1.0 vs 1.2.0 support [see FAQ](https://github.com/PyTorchLightning/pytorch-lightning#faq)    
-->

## 🐛 Bug

<!-- A clear and concise description of what the bug is. -->

### To Reproduce

I use dgl library to make a gnn and batch the DGLGraph. 
No problem during training, but in test, I got a TypeError: to() got an unexpected keyword argument 'non_blocking'

<class 'dgl.graph.DGLGraph'> .to() function has no keyword argument 'non_blocking'

<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->


#### Code sample
<!-- Ideally attach a minimal code sample to reproduce the decried issue. 
Minimal means having the shortest code but still preserving the bug. -->

### Expected behavior

<!-- A clear and concise description of what you expected to happen. -->

### Environment


 - OS: Linux
 - CUDA: 10.1
 - Python Version: 3.7
 - PyTorch Version: 1.5.1
 - DGL Version: 0.4.3post2
 - PyTorch-Lightning Version: 0.8.5

### Additional context

```
   File "../src/main.py", line 131, in <module>
    run(params)
  File "../src/main.py", line 92, in run
    trainer.test(model)
  File "/home/jiangyize/miniconda3/envs/galixir/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 1279, in test
    results = self.__test_given_model(model, test_dataloaders)
  File "/home/jiangyize/miniconda3/envs/galixir/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 1346, in __test_given_model
    results = self.fit(model)
  File "/home/jiangyize/miniconda3/envs/galixir/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 1003, in fit
    results = self.single_gpu_train(model)
  File "/home/jiangyize/miniconda3/envs/galixir/lib/python3.7/site-packages/pytorch_lightning/trainer/distrib_parts.py", line 186, in single_gpu_train
    results = self.run_pretrain_routine(model)
  File "/home/jiangyize/miniconda3/envs/galixir/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 1166, in run_pretrain_routine
    results = self.run_evaluation(test_mode=True)
  File "/home/jiangyize/miniconda3/envs/galixir/lib/python3.7/site-packages/pytorch_lightning/trainer/evaluation_loop.py", line 391, in run_evaluation
    eval_results = self._evaluate(self.model, dataloaders, max_batches, test_mode)
  File "/home/jiangyize/miniconda3/envs/galixir/lib/python3.7/site-packages/pytorch_lightning/trainer/evaluation_loop.py", line 293, in _evaluate
    output = self.evaluation_forward(model, batch, batch_idx, dataloader_idx, test_mode)
  File "/home/jiangyize/miniconda3/envs/galixir/lib/python3.7/site-packages/pytorch_lightning/trainer/evaluation_loop.py", line 458, in evaluation_forward
    batch = self.transfer_batch_to_gpu(batch, root_gpu)
  File "/home/jiangyize/miniconda3/envs/galixir/lib/python3.7/site-packages/pytorch_lightning/trainer/distrib_parts.py", line 159, in transfer_batch_to_gpu
    return self.__transfer_batch_to_device(batch, device)
  File "/home/jiangyize/miniconda3/envs/galixir/lib/python3.7/site-packages/pytorch_lightning/trainer/distrib_parts.py", line 164, in __transfer_batch_to_device
    return model.transfer_batch_to_device(batch, device)
  File "/home/jiangyize/miniconda3/envs/galixir/lib/python3.7/site-packages/pytorch_lightning/core/hooks.py", line 242, in transfer_batch_to_device
    return move_data_to_device(batch, device)
  File "/home/jiangyize/miniconda3/envs/galixir/lib/python3.7/site-packages/pytorch_lightning/utilities/apply_func.py", line 109, in move_data_to_device
    return apply_to_collection(batch, dtype=(TransferableDataType, Batch), function=batch_to)
  File "/home/jiangyize/miniconda3/envs/galixir/lib/python3.7/site-packages/pytorch_lightning/utilities/apply_func.py", line 40, in apply_to_collection
    for k, v in data.items()})
  File "/home/jiangyize/miniconda3/envs/galixir/lib/python3.7/site-packages/pytorch_lightning/utilities/apply_func.py", line 40, in <dictcomp>
    for k, v in data.items()})
  File "/home/jiangyize/miniconda3/envs/galixir/lib/python3.7/site-packages/pytorch_lightning/utilities/apply_func.py", line 35, in apply_to_collection
    return function(data, *args, **kwargs)
  File "/home/jiangyize/miniconda3/envs/galixir/lib/python3.7/site-packages/pytorch_lightning/utilities/apply_func.py", line 107, in batch_to
    return data.to(device, non_blocking=True)
TypeError: to() got an unexpected keyword argument 'non_blocking'
```


Discussion:
Hi! thanks for your contribution!, great first issue!
Having the same problem; it's because `DGLGraph.to` ([docs](https://docs.dgl.ai/en/0.4.x/generated/dgl.DGLGraph.to.html#dgl.DGLGraph.to), [source](https://docs.dgl.ai/en/0.4.x/_modules/dgl/graph.html#DGLGraph.to)) doesn't take the non_blocking argument. Example:

```python
dgl.DGLGraph().to('cuda', non_blocking=True)
```

Here's my temporary solution:

```python
class LightningDGLGraph(DGLGraph):
    def to(self, ctx, *args, **kwargs):
        return super().to(torch.device(ctx))

g = LightningDGLGraph()
g.to('cuda', non_blocking=True)
```

Works, but probably not ideal.
> Having the same problem; it's because `DGLGraph.to` ([docs](https://docs.dgl.ai/en/0.4.x/generated/dgl.DGLGraph.to.html#dgl.DGLGraph.to), [source](https://docs.dgl.ai/en/0.4.x/_modules/dgl/graph.html#DGLGraph.to)) doesn't take the non_blocking argument. Example:
> 
> ```python
> dgl.DGLGraph().to('cuda', non_blocking=True)
> ```
> 
> Here's my temporary solution:
> 
> ```python
> class LightningDGLGraph(DGLGraph):
>     def to(self, ctx, *args, **kwargs):
>         return super().to(torch.device(ctx))
> 
> g = LightningDGLGraph()
> g.to('cuda', non_blocking=True)
> ```
> 
> Works, but probably not ideal.

Hi, I wonder how this will work if using dgl.batch? the class type they return to you is a `DGLGraph`.


Ok, they also have this quick fix here: https://github.com/dmlc/dgl/pull/1600.
so uninstall the stable version and install the latest version from main solves my problem:
```
pip install --pre dgl           # For CPU Build
pip install --pre dgl-cu90      # For CUDA 9.0 Build
pip install --pre dgl-cu92      # For CUDA 9.2 Build
pip install --pre dgl-cu100     # For CUDA 10.0 Build
pip install --pre dgl-cu101     # For CUDA 10.1 Build
```

It seems to work for me for now.
> Ok, they also have this quick fix here: [dmlc/dgl#1600](https://github.com/dmlc/dgl/pull/1600).
> so uninstall the stable version and install the latest version from main solves my problem:

Just saw your edit. ~~This seems to work if I **don't** specify the number of gpus; when I do, same error.~~ E: It's the distributed backend; it never calls `graph.to`. You can throw a `0/0` in there and it'll never break with `distributed_backend ddp`.
For a clean solution in Lightning, override [this](https://pytorch-lightning.readthedocs.io/en/latest/api/pytorch_lightning.core.hooks.html?highlight=transfer#pytorch_lightning.core.hooks.ModelHooks.transfer_batch_to_device) model hook and call .to() yourself on the graph object.

Regarding ddp, is DGLGraph supposed to work with that (I mean in plain pytorch)? I don't think it can work with scatter and gather.
> For a clean solution in Lightning, override this model hook and call .to() yourself on the graph object.

@awaelchli Is this supposed to be overridden in the model? It doesn't seem to get called for me in a distributed setting.
@jacobdanovitch Yes, this hook only works for single gpu, because in distributed we need to scatter and gather a batch, and if it is a custom object we don't know how to do that. For this, you would have to define your own DistributedDataParallel module and configure it in the [configure_ddp ](https://pytorch-lightning.readthedocs.io/en/latest/api/pytorch_lightning.core.html#pytorch_lightning.core.LightningModule.configure_ddp)model hook. We should probably update the docs regarding that.

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
@@ -309,8 +309,13 @@ def transfer_batch_to_device(self, batch, device)
         Note:
             This hook should only transfer the data and not modify it, nor should it move the data to
             any other device than the one passed in as argument (unless you know what you are doing).
-            The :class:`~pytorch_lightning.trainer.trainer.Trainer` already takes care of splitting the
-            batch and determines the target devices.
+
+        Note:
+            This hook only runs on single GPU training (no data-parallel). If you need multi-GPU support
+            for your custom batch objects, you need to define your custom
+            :class:`~torch.nn.parallel.DistributedDataParallel` or
+            :class:`~pytorch_lightning.overrides.data_parallel.LightningDistributedDataParallel` and
+            override :meth:`~pytorch_lightning.core.lightning.LightningModule.configure_ddp`.
 
         See Also:
             - :func:`~pytorch_lightning.utilities.apply_func.move_data_to_device`
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
@@ -104,6 +104,7 @@ def batch_to(data):
                 setattr(device_data, field, device_field)
             return device_data
 
-        return data.to(device, non_blocking=True)
+        kwargs = dict(non_blocking=True) if isinstance(data, torch.Tensor) else {}
+        return data.to(device, **kwargs)
 
     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
@@ -1,4 +1,5 @@
 from collections import namedtuple
+from unittest.mock import patch
 
 import pytest
 import torch
@@ -384,3 +385,24 @@ def to(self, *args, **kwargs):
 
     assert batch.text.type() == 'torch.cuda.LongTensor'
     assert batch.label.type() == 'torch.cuda.LongTensor'
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="test requires GPU machine")
+def test_non_blocking():
+    """ Tests that non_blocking=True only gets passed on torch.Tensor.to, but not on other objects. """
+    trainer = Trainer()
+
+    batch = torch.zeros(2, 3)
+    with patch.object(batch, 'to', wraps=batch.to) as mocked:
+        trainer.transfer_batch_to_gpu(batch, 0)
+        mocked.assert_called_with(torch.device('cuda', 0), non_blocking=True)
+
+    class BatchObject(object):
+
+        def to(self, *args, **kwargs):
+            pass
+
+    batch = BatchObject()
+    with patch.object(batch, 'to', wraps=batch.to) as mocked:
+        trainer.transfer_batch_to_gpu(batch, 0)
+        mocked.assert_called_with(torch.device('cuda', 0))