updated to_device function with auto and None arguments to:

    *  use the fastest device available
    *  leave the device type unmodified respectively
(skorch-dev#600)

Author: BAjayiobe1

CHANGES.md

@@ -11,6 +11,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Added `NeptuneLogger` callback for logging experiment metadata to neptune.ai
 - Add DataFrameTransformer, an sklearn compatible transformer that helps working with pandas DataFrames by transforming the DataFrame into a representation that works well with neural networks (#507)
+- Added `auto` and `None` options to automatically use the fastest device available or to leave the device type unmodified respectively. (#600)
 
 ### Changed
 

examples/benchmarks/mnist.py

@@ -34,6 +34,7 @@
 from sklearn.model_selection import StratifiedKFold
 from sklearn.metrics import accuracy_score
 from sklearn.utils import shuffle
+from skorch.utils import to_device
 from skorch import NeuralNetClassifier
 from skorch.callbacks import EpochScoring
 import torch
@@ -141,7 +142,7 @@ def train_torch(
         lr,
         max_epochs,
 ):
-    model.to(device)
+    model = to_device(model, device)
 
     idx_train, idx_valid = next(iter(StratifiedKFold(
         5, random_state=0).split(np.arange(len(X)), y)))
@@ -191,7 +192,7 @@ def train_step(model, dataset, device, criterion, batch_size, optimizer):
     batch_sizes = []
     tic = time.time()
     for Xi, yi in torch.utils.data.DataLoader(dataset, batch_size=batch_size):
-        Xi, yi = Xi.to(device), yi.to(device)
+        Xi, yi = to_device(Xi, device), to_device(yi, device)
         optimizer.zero_grad()
         y_pred = model(Xi)
         y_pred = torch.log(y_pred)
@@ -221,7 +222,7 @@ def valid_step(model, dataset, device, criterion, batch_size):
         for Xi, yi in torch.utils.data.DataLoader(
                 dataset, batch_size=batch_size,
         ):
-            Xi, yi = Xi.to(device), yi.to(device)
+            Xi, yi = to_device(Xi, device), to_device(yi, device)
             y_pred = model(Xi)
             y_pred = torch.log(y_pred)
             loss = criterion(y_pred, yi)

examples/word_language_model/data.py

@@ -1,5 +1,7 @@
 import os
 
+from skorch.utils import to_device
+
 import torch
 from torch.autograd import Variable
 
@@ -70,7 +72,7 @@ def batchify(self, data, bsz):
         data = data.narrow(0, 0, nbatch * bsz)
         # Evenly divide the data across the bsz batches.
         data = data.view(bsz, -1).t().contiguous()
-        return data.to(self.device)
+        return to_device(data, self.device)
 
     def get_batch(self, i):
         seq_len = min(self.bptt, len(self.batches) - 1 - i)

skorch/net.py

@@ -431,7 +431,7 @@ def initialize_criterion(self):
         criterion_params = self._get_params_for('criterion')
         self.criterion_ = self.criterion(**criterion_params)
         if isinstance(self.criterion_, torch.nn.Module):
-            self.criterion_ = self.criterion_.to(self.device)
+            self.criterion_ = to_device(self.criterion_, self.device)
         return self
 
     def _format_reinit_msg(self, name, kwargs=None, triggered_directly=True):
@@ -472,7 +472,7 @@ def initialize_module(self):
 
             module = module(**kwargs)
 
-        self.module_ = module.to(self.device)
+        self.module_ = to_device(module, self.device)
         return self
 
     def _is_virtual_param(self, key):

skorch/tests/test_utils.py

@@ -149,40 +149,94 @@ def x(self):
     def x_tup(self):
         return torch.zeros(3), torch.ones((4, 5))
 
+    @pytest.fixture
+    def x_pad_seq(self):
+        input = torch.zeros((5, 3)).float()
+        length = torch.as_tensor([2, 2, 1])
+        return pack_padded_sequence(input,length )
+
+    def check_device_type(self, tensor, device_input, prev_device):
+        """assert expected device type conditioned on the input argument for `to_device`"""
+        if None is device_input:
+            assert tensor.device.type == prev_device
+
+        elif device_input == "auto":
+            expected = 'cuda' if torch.cuda.is_available() else 'cpu'
+            assert tensor.device.type == expected
+
+        else:
+            assert tensor.device.type == device_input
+
+
     @pytest.mark.parametrize('device_from, device_to', [
         ('cpu', 'cpu'),
         ('cpu', 'cuda'),
         ('cuda', 'cpu'),
         ('cuda', 'cuda'),
+        (None, None),
+        ('auto', 'auto'),
     ])
     def test_check_device_torch_tensor(self, to_device, x, device_from, device_to):
         if 'cuda' in (device_from, device_to) and not torch.cuda.is_available():
             pytest.skip()
 
+        prev_device = None
+        if None in (device_from, device_to):
+            prev_device = x.device.type
+
         x = to_device(x, device=device_from)
-        assert x.device.type == device_from
+        self.check_device_type(x, device_from, prev_device)
 
         x = to_device(x, device=device_to)
-        assert x.device.type == device_to
+        self.check_device_type(x, device_to, prev_device)
 
     @pytest.mark.parametrize('device_from, device_to', [
         ('cpu', 'cpu'),
         ('cpu', 'cuda'),
         ('cuda', 'cpu'),
         ('cuda', 'cuda'),
+        (None, None),
+        ('auto', 'auto'),
     ])
     def test_check_device_tuple_torch_tensor(
-            self, to_device, x, device_from, device_to):
+            self, to_device, x_tup, device_from, device_to):
         if 'cuda' in (device_from, device_to) and not torch.cuda.is_available():
             pytest.skip()
 
-        x = to_device(x, device=device_from)
-        for xi in x:
-            assert xi.device.type == device_from
+        prev_device = [None for _ in range(len(x_tup))]
+        if None in (device_from, device_to):
+            prev_device = [x.device.type for x in x_tup]
 
-        x = to_device(x, device=device_to)
-        for xi in x:
-            assert xi.device.type == device_to
+        x_tup = to_device(x_tup, device=device_from)
+        for idx, xi in enumerate(x_tup):
+            self.check_device_type(xi, device_from, prev_device[idx])
+
+        x_tup = to_device(x_tup, device=device_to)
+        for idx, xi in enumerate(x_tup):
+            self.check_device_type(xi, device_to, prev_device[idx])
+
+    @pytest.mark.parametrize('device_from, device_to', [
+        ('cpu', 'cpu'),
+        ('cpu', 'cuda'),
+        ('cuda', 'cpu'),
+        ('cuda', 'cuda'),
+        (None, None),
+        ('auto', 'auto'),
+    ])
+    def test_check_device_packed_padded_sequence(
+            self, to_device, x_pad_seq, device_from, device_to):
+        if 'cuda' in (device_from, device_to) and not torch.cuda.is_available():
+            pytest.skip()
+
+        prev_device = None
+        if None in (device_from, device_to):
+            prev_device = x_pad_seq.data.device.type
+
+        x_pad_seq = to_device(x_pad_seq, device=device_from)
+        self.check_device_type(x_pad_seq.data, device_from, prev_device)
+
+        x_pad_seq = to_device(x_pad_seq, device=device_to)
+        self.check_device_type(x_pad_seq.data, device_to, prev_device)
 
 
 class TestDuplicateItems:

skorch/utils.py

@@ -80,7 +80,7 @@ def to_tensor(X, device, accept_sparse=False):
     to_tensor_ = partial(to_tensor, device=device)
 
     if is_torch_data_type(X):
-        return X.to(device)
+        return to_device(X, device)
     if isinstance(X, dict):
         return {key: to_tensor_(val) for key, val in X.items()}
     if isinstance(X, (list, tuple)):
@@ -126,12 +126,33 @@ def to_numpy(X):
 
 
 def to_device(X, device):
-    """Generic function to move module output(s) to a device.
+    """Generic function to modify the device type of the tensor(s)/module inputted.
 
-    Deals with X being a torch tensor or a tuple of torch tensors.
+    Parameters
+    ----------
+    X : input data
+        Deals with X being a:
+
+         * torch tensor
+         * tuple of torch tensors
+         * PackSequence instance
+         * torch.nn.Module
+
+    device : str, torch.device
+        The compute device to be used. If device="auto" it is set to
+        "cuda" if available otherwise "cpu". If device=None, return
+        the input unmodified
 
     """
-    if isinstance(X, tuple):
+    if device is None:
+        return X
+
+    if device == "auto":
+        use_cuda = torch.cuda.is_available()
+        device = "cuda" if use_cuda else "cpu"
+
+    # PackedSequence class inherits from a namedtuple
+    if isinstance(X, tuple) & (type(X) != PackedSequence):
         return tuple(x.to(device) for x in X)
     return X.to(device)