From ec74e1aa5f46f684ba64b481040584b0d9d75ee0 Mon Sep 17 00:00:00 2001
From: Abe Levitan <alevitan@mit.edu>
Date: Mon, 1 Apr 2019 09:31:33 -0400
Subject: [PATCH] Add the ability for datasets to change what device they
 output to even if data remains stored on cpu

---
 CDTools/datasets.py    | 33 ++++++++++++++++++++++++++++++++-
 tests/test_datasets.py | 25 ++++++++++++++++++++++++-
 2 files changed, 56 insertions(+), 2 deletions(-)

diff --git a/CDTools/datasets.py b/CDTools/datasets.py
index c2e83f3..6fc77b3 100644
--- a/CDTools/datasets.py
+++ b/CDTools/datasets.py
@@ -107,6 +107,11 @@ class CDataset(torchdata.Dataset):
         else:
             self.mask = None
 
+        if t.cuda.is_available():
+            self.get_as(device='cuda:0')
+        else:
+            self.get_as(device='cpu')
+
             
     def to(self,*args,**kwargs):
         # The mask should always stay a uint8, but it should switch devices
@@ -119,7 +124,33 @@ class CDataset(torchdata.Dataset):
         if self.mask is not None:
             self.mask = self.mask.to(*args,**mask_kwargs)          
 
+
+    def get_as(self, *args, **kwargs):
+        self.get_as_args = (args, kwargs)
+
+
+    def __getitem__(self, index):
+        # Deals with loading to appropriate device/dtype, if
+        # specified via a call to get_as
+        inputs, outputs = self._load(index)
+        if hasattr(self, 'get_as_args'):
+            outputs = outputs.to(*self.get_as_args[0],**self.get_as_args[1])
+            moved_inputs = []
+            for inp in inputs:
+                try:
+                    moved_inputs.append(inp.to(*self.get_as_args[0],**self.get_as_args[1]) )
+                except:
+                    moved_inputs.append(inp)
+        else:
+            moved_inputs = inputs
+        return moved_inputs, outputs
+
     
+    def _load(self, index):
+        # Internal function to load data
+        raise NotImplementedError()
+        
+            
     @classmethod
     def from_cxi(cls, cxi_file):
         entry_info = cdtdata.get_entry_info(cxi_file)
@@ -177,7 +208,7 @@ class Ptycho_2D_Dataset(CDataset):
     def __len__(self):
         return self.patterns.shape[0]
 
-    def __getitem__(self, index):
+    def _load(self, index):
         return (index, self.translations[index]), self.patterns[index]
 
 
diff --git a/tests/test_datasets.py b/tests/test_datasets.py
index 52620c7..aa0d6d8 100644
--- a/tests/test_datasets.py
+++ b/tests/test_datasets.py
@@ -27,6 +27,7 @@ def test_CDataset_init():
     mask = np.ones((256,256))
     dataset = CDataset(entry_info, sample_info,
                        wavelength, detector_geometry, mask)
+    
     assert t.all(t.eq(dataset.mask,t.tensor(mask)))
     assert dataset.entry_info == entry_info
     assert dataset.sample_info == sample_info
@@ -231,13 +232,35 @@ def test_Ptycho_2D_Dataset_to(ptycho_cxi_1):
         assert dataset.patterns.device == t.device('cuda:0')
         assert dataset.translations.device == t.device('cuda:0')
 
-    
+
+        
 def test_Ptycho_2D_Dataset_ops(ptycho_cxi_1):
     cxi, expected = ptycho_cxi_1
     dataset = Ptycho_2D_Dataset.from_cxi(cxi)
+    dataset.get_as('cpu')
 
     assert len(dataset) == expected['data'].shape[0]
     (idx, translation), pattern = dataset[3]
     assert idx == 3
     assert t.allclose(translation, t.tensor(expected['translations'][3,:]))
     assert t.allclose(pattern, t.tensor(expected['data'][3,:,:]))
+
+
+def test_Ptycho_2D_Dataset_get_as(ptycho_cxi_1):
+    cxi, expected = ptycho_cxi_1
+    dataset = Ptycho_2D_Dataset.from_cxi(cxi)
+    if t.cuda.is_available():
+        dataset.get_as('cuda:0')
+        assert len(dataset) == expected['data'].shape[0]
+
+        (idx, translation), pattern = dataset[3]
+        assert str(translation.device) == 'cuda:0'
+        assert str(pattern.device) == 'cuda:0'
+        
+        assert idx == 3
+        assert t.allclose(translation.to(device='cpu'),
+                          t.tensor(expected['translations'][3,:]))
+        assert t.allclose(pattern.to(device='cpu'),
+                          t.tensor(expected['data'][3,:,:]))
+        
+