Repository: ElementAI/baal (python)

Issue:
Support tuple and dicts for *_on_batch
**Describe the bug**
If the dataset returns a Dict or Tuple, it breaks.


**To Reproduce**
```python
class MyDS(Dataset):
	def __len__(self):
		return 5

	def __getitem__(self,idx):
		return (torch.randn(3, 20, 20), torch.randn(3, 20, 20)), 5
...

wrapper.train_on_batch(MyDS()[1], ...)

**Expected behavior**

The list of tensors should be acceted and `cuda` should be called on them.

**Vesion (please complete the following information):**
 - BaaL version: 1.0


Pull request description:
Flexible cuda calls
## Summary:

Move data to the GPU in a more flexible manner. This closes #17 as it will work for dictionaries, tuples, and basically anything we can expect to get out of a dataloader.

### Features:

`baal.utils.cuda_utils.to_cuda` is a function that will move a tensor, tensors in any "sequence"-container, and tensors that are values in dictionaries, to the GPU.

## Checklist:

* [x] Your code is documented (To validate this, add your module to `tests/documentation_test.py`).
* [x] Your code is tested with unit tests.


Patch:
diff --git a/.coveragerc b/.coveragerc
index 718f968..547ae5e 100644
--- a/.coveragerc
+++ b/.coveragerc
@@ -7,6 +7,8 @@ exclude_lines =
     raise NotImplementedError
     abstractmethod
     __repr__
+    # exclude lines where cuda is called on anything:
+    cuda\(
 
 [run]
 source =
diff --git a/src/baal/modelwrapper.py b/src/baal/modelwrapper.py
index 87aecb6..dd7c24a 100644
--- a/src/baal/modelwrapper.py
+++ b/src/baal/modelwrapper.py
@@ -6,13 +6,13 @@ from typing import Callable, Optional
 import numpy as np
 import structlog
 import torch
-from torch.autograd import Variable
 from torch.optim import Optimizer
 from torch.utils.data import DataLoader, Dataset
 from torch.utils.data.dataloader import default_collate
 from tqdm import tqdm
 
 from baal.utils.iterutils import map_on_tensor
+from baal.utils.cuda_utils import to_cuda
 from baal.utils.metrics import Loss
 
 log = structlog.get_logger("ModelWrapper")
@@ -279,9 +279,9 @@ class ModelWrapper:
         Returns:
             Tensor, the loss computed from the criterion.
         """
-        data = Variable(data)
+
         if cuda:
-            data, target = data.cuda(), target.cuda()
+            data, target = to_cuda(data), to_cuda(target)
         optimizer.zero_grad()
         output = self.model(data)
         loss = self.criterion(output, target)
@@ -312,7 +312,7 @@ class ModelWrapper:
         """
         with torch.no_grad():
             if cuda:
-                data, target = data.cuda(), target.cuda()
+                data, target = to_cuda(data), to_cuda(target)
             if average_predictions == 1:
                 preds = self.model(data)
                 loss = self.criterion(preds, target)
@@ -343,7 +343,7 @@ class ModelWrapper:
         """
         with torch.no_grad():
             if cuda:
-                data = data.cuda()
+                data = to_cuda(data)
             if self.replicate_in_memory:
                 input_shape = data.size()
                 batch_size = input_shape[0]
diff --git a/src/baal/utils/cuda_utils.py b/src/baal/utils/cuda_utils.py
new file mode 100644
index 0000000..167de17
--- /dev/null
+++ b/src/baal/utils/cuda_utils.py
@@ -0,0 +1,47 @@
+from functools import singledispatch
+from collections.abc import Mapping, Sequence
+import torch
+
+
+@singledispatch
+def to_cuda(data):
+    """
+    Move an object to CUDA.
+
+    This function works recursively on lists and dicts, moving the values
+    inside to cuda.
+
+    Args:
+        data (list, tuple, dict, torch.Tensor, torch.nn.Module):
+            The data you'd like to move to the GPU. If there's a pytorch tensor or
+            model in data (e.g. in a list or as values in a dictionary) this
+            function will move them all to CUDA and return something that matches
+            the input in structure.
+
+    Returns:
+        list, tuple, dict, torch.Tensor, torch.nn.Module:
+            Data of the same type / structure as the input.
+    """
+    # the base case: if this is not a type we recognise, return it
+    return data
+
+
+@to_cuda.register(torch.Tensor)
+@to_cuda.register(torch.nn.Module)
+def _(data):
+    return data.cuda()
+
+
+@to_cuda.register
+def _(data: Mapping):
+    # use the type of the object to create a new one:
+    return type(data)([(key, to_cuda(val)) for key, val in data.items()])
+
+
+@to_cuda.register
+def _(data: Sequence):
+    # use the type of this object to instantiate a new one:
+    if hasattr(data, "_fields"):  # in case it's a named tuple
+        return type(data)(*(to_cuda(item) for item in data))
+    else:
+        return type(data)(to_cuda(item) for item in data)


Test patch:
diff --git a/tests/documentation_test.py b/tests/documentation_test.py
index 0874d20..885a105 100644
--- a/tests/documentation_test.py
+++ b/tests/documentation_test.py
@@ -6,7 +6,7 @@ from itertools import compress
 
 import pytest
 
-modules = ['baal.active', 'baal', 'baal.utils.metrics']
+modules = ['baal.active', 'baal', 'baal.utils.metrics', 'baal.utils.cuda_utils']
 # We do not force these functions to be compliant.
 accepted_name = ['set_logger_config']
 # We do not force these modules to be compliant.
diff --git a/tests/utils/test_cuda_utils.py b/tests/utils/test_cuda_utils.py
new file mode 100644
index 0000000..dfd9d7f
--- /dev/null
+++ b/tests/utils/test_cuda_utils.py
@@ -0,0 +1,36 @@
+import os
+import pytest
+import torch
+from collections import namedtuple, OrderedDict
+
+from baal.utils.cuda_utils import to_cuda
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="No CUDA available")
+def test_to_cuda():
+    t = torch.randn(3, 5)
+    assert to_cuda(t).device.type == "cuda"
+    m = torch.nn.Linear(3, 1)
+    assert next(to_cuda(m).parameters()).device.type == "cuda"
+
+    t_list = [t, t.clone()]
+    assert type(t_list) is type(to_cuda(t_list))
+    assert all(t_.device.type == "cuda" for t_ in to_cuda(t_list))
+
+    t_dict = {"t": t, "t2": t.clone()}
+    assert type(t_dict) is type(to_cuda(t_dict))
+    assert all(t_.device.type == "cuda" for t_ in to_cuda(t_dict).values())
+
+    MyTuple = namedtuple("MyTuple", ["t", "t2"])
+    t_named_tuple = MyTuple(t, t.clone())
+    assert type(t_named_tuple) is type(to_cuda(t_named_tuple))
+    assert all(t_.device.type == "cuda" for t_ in to_cuda(t_named_tuple))
+
+    t_tuple = (t, t.clone())
+    assert type(t_tuple) is type(to_cuda(t_tuple))
+    assert all(t_.device.type == "cuda" for t_ in to_cuda(t_tuple))
+
+    # test a type that's not explicitly defined:
+    t_ordered_dict = OrderedDict([("t", t), ("t2", t.clone())])
+    assert type(t_ordered_dict) is type(to_cuda(t_ordered_dict))
+    assert all(t_.device.type == "cuda" for t_ in to_cuda(t_ordered_dict).values())