- integrated device transfers in graph model

inferno/extensions/containers/graph.py

@@ -10,6 +10,7 @@
 
 from ...utils import python_utils as pyu
 from ...utils.exceptions import assert_
+from ..layers.device import OnDevice
 
 
 __all__ = ['NNGraph', 'Graph']
@@ -360,6 +361,21 @@ def get_module_for_nodes(self, names):
             modules.append(module)
         return pyu.from_iterable(modules)
 
+    def to_device(self, names, target_device, device_ordinal=None, async=False):
+        """Transfer nodes in the network to a specified device."""
+        names = pyu.to_iterable(names)
+        for name in names:
+            assert self.is_node_in_graph(name), "Node '{}' is not in graph.".format(name)
+            module = getattr(self, name, None)
+            assert module is not None, "Node '{}' is in the graph but could not find a module " \
+                                       "corresponding to it.".format(name)
+            # Transfer
+            module_on_device = OnDevice(module, target_device,
+                                        device_ordinal=device_ordinal,
+                                        async=async)
+            setattr(self, name, module_on_device)
+        return self
+
     def get_parameters_for_nodes(self, names, named=False):
         """Get parameters of all nodes listed in `names`."""
         if not named:

inferno/extensions/layers/device.py

@@ -2,7 +2,7 @@
 from ...utils.python_utils import from_iterable, to_iterable
 from ...utils.exceptions import assert_, DeviceError
 
-__all__ = ['DeviceTransfer']
+__all__ = ['DeviceTransfer', 'OnDevice']
 
 
 class DeviceTransfer(nn.Module):

tests/extensions/containers/graph.py

@@ -99,6 +99,23 @@ def test_graph_basic(self):
         model.add_output_node('output_0', previous='conv1')
         ModelTester((1, 1, 100, 100), (1, 1, 100, 100))(model)
 
+    def test_graph_device_transfers(self):
+        from inferno.extensions.containers.graph import Graph
+        from inferno.extensions.layers.convolutional import ConvELU2D
+        import torch
+        from torch.autograd import Variable
+        # Build graph
+        model = Graph()
+        model.add_input_node('input_0')
+        model.add_node('conv0', ConvELU2D(1, 10, 3), previous='input_0')
+        model.add_node('conv1', ConvELU2D(10, 1, 3), previous='conv0')
+        model.add_output_node('output_0', previous='conv1')
+        # Transfer
+        model.to_device('conv0', 'cpu').to_device('conv1', 'cuda', 0)
+        x = Variable(torch.rand(1, 1, 100, 100))
+        y = model(x)
+        self.assertIsInstance(y.data, torch.cuda.FloatTensor)
+
     @unittest.skip("Needs machine with 4 GPUs")
     def test_multi_gpu(self):
         import torch