From db9254acbe61fbba97da5f9e886d4bac284b1b3f Mon Sep 17 00:00:00 2001
From: William Falcon <waf2107@columbia.edu>
Date: Thu, 15 Aug 2019 09:39:09 -0400
Subject: [PATCH] enable recursive parsing for single gpu inputs (#121)

* added tests

* added single gpu data transfer recursive

* added single gpu data transfer recursive

* added single gpu data transfer recursive

* added single gpu data transfer recursive

* added single gpu data transfer recursive

* added single gpu data transfer recursive
---
 pytorch_lightning/models/trainer.py | 28 ++++++++++++++++++-----
 tests/test_models.py                | 35 +++++++++++++++++++++++++++++
 2 files changed, 57 insertions(+), 6 deletions(-)

diff --git a/pytorch_lightning/models/trainer.py b/pytorch_lightning/models/trainer.py
index cc3c82f6..6d4e4e3f 100644
--- a/pytorch_lightning/models/trainer.py
+++ b/pytorch_lightning/models/trainer.py
@@ -390,9 +390,8 @@ class Trainer(TrainerIO):
         elif self.single_gpu:
             # put inputs on gpu manually
             gpu_id = self.data_parallel_device_ids[0]
-            for i, x in enumerate(data_batch):
-                if isinstance(x, torch.Tensor):
-                    data_batch[i] = x.cuda(gpu_id)
+            data_batch = self.transfer_batch_to_gpu(data_batch, gpu_id)
+            args[0] = data_batch
 
             # do non dp, ddp step
             output = model.validation_step(*args)
@@ -905,6 +904,24 @@ We recommend you switch to ddp if you want to use amp
         blacklist = {'batch_nb', 'v_nb', 'gpu'}
         return blacklist
 
+    def transfer_batch_to_gpu(self, batch, gpu_id):
+        # base case
+        if isinstance(batch, torch.Tensor):
+            return batch.cuda(gpu_id)
+
+        # when list
+        elif isinstance(batch, list):
+            for i, x in enumerate(batch):
+                batch[i] = self.transfer_batch_to_gpu(x, gpu_id)
+            return batch
+
+        # when dict
+        elif isinstance(batch, dict):
+            for k, v in batch.items():
+                batch[k] = self.transfer_batch_to_gpu(v, gpu_id)
+
+            return batch
+
     def __tng_forward(self, data_batch, batch_nb, opt_idx):
         """
         Handle forward for each training case (distributed, single gpu, etc...)
@@ -926,9 +943,8 @@ We recommend you switch to ddp if you want to use amp
             output = self.model(*args)
         elif self.single_gpu:
             gpu_id = self.data_parallel_device_ids[0]
-            for i, x in enumerate(data_batch):
-                if isinstance(x, torch.Tensor):
-                    data_batch[i] = x.cuda(gpu_id)
+            data_batch = self.transfer_batch_to_gpu(data_batch, gpu_id)
+            args[0] = data_batch
             output = self.model.training_step(*args)
 
         else:
diff --git a/tests/test_models.py b/tests/test_models.py
index 992bc7a8..b7722584 100644
--- a/tests/test_models.py
+++ b/tests/test_models.py
@@ -26,6 +26,41 @@ np.random.seed(SEED)
 # ------------------------------------------------------------------------
 # TESTS
 # ------------------------------------------------------------------------
+def test_single_gpu_batch_parse():
+    if not torch.cuda.is_available():
+        warnings.warn('test_amp_gpu_ddp cannot run.'
+                      'Rerun on a GPU node to run this test')
+        return
+    if not torch.cuda.device_count() > 1:
+        warnings.warn('test_amp_gpu_ddp cannot run.'
+                      'Rerun on a node with 2+ GPUs to run this test')
+        return
+
+    trainer = Trainer()
+
+    # batch is just a tensor
+    batch = torch.rand(2, 3)
+    batch = trainer.transfer_batch_to_gpu(batch, 0)
+    assert batch.device.index == 0 and batch.type() == 'torch.cuda.FloatTensor'
+
+    # tensor list
+    batch = [torch.rand(2, 3), torch.rand(2, 3)]
+    batch = trainer.transfer_batch_to_gpu(batch, 0)
+    assert batch[0].device.index == 0 and batch[0].type() == 'torch.cuda.FloatTensor'
+    assert batch[1].device.index == 0 and batch[1].type() == 'torch.cuda.FloatTensor'
+
+    # tensor list of lists
+    batch = [[torch.rand(2, 3), torch.rand(2, 3)]]
+    batch = trainer.transfer_batch_to_gpu(batch, 0)
+    assert batch[0][0].device.index == 0 and batch[0][0].type() == 'torch.cuda.FloatTensor'
+    assert batch[0][1].device.index == 0 and batch[0][1].type() == 'torch.cuda.FloatTensor'
+
+    # tensor dict
+    batch = [{'a': torch.rand(2, 3), 'b': torch.rand(2, 3)}]
+    batch = trainer.transfer_batch_to_gpu(batch, 0)
+    assert batch[0]['a'].device.index == 0 and batch[0]['a'].type() == 'torch.cuda.FloatTensor'
+    assert batch[0]['b'].device.index == 0 and batch[0]['b'].type() == 'torch.cuda.FloatTensor'
+
 
 def test_early_stopping_cpu_model():
     """