Improvement in DDP and bug fix in DataModule (#432)

pina/data/dataset.py

@@ -6,6 +6,7 @@ from torch.utils.data import Dataset
 from abc import abstractmethod
 from torch_geometric.data import Batch
 
+
 class PinaDatasetFactory:
     """
     Factory class for the PINA dataset. Depending on the type inside the
@@ -13,6 +14,7 @@ class PinaDatasetFactory:
     - PinaTensorDataset for torch.Tensor
     - PinaGraphDataset for list of torch_geometric.data.Data objects
     """
+
     def __new__(cls, conditions_dict, **kwargs):
         if len(conditions_dict) == 0:
             raise ValueError('No conditions provided')
@@ -25,10 +27,12 @@ class PinaDatasetFactory:
         raise ValueError('Conditions must be either torch.Tensor or list of Data '
                          'objects.')
 
+
 class PinaDataset(Dataset):
     """
     Abstract class for the PINA dataset
     """
+
     def __init__(self, conditions_dict, max_conditions_lengths):
         self.conditions_dict = conditions_dict
         self.max_conditions_lengths = max_conditions_lengths
@@ -49,6 +53,7 @@ class PinaDataset(Dataset):
     def __getitem__(self, item):
         pass
 
+
 class PinaTensorDataset(PinaDataset):
     def __init__(self, conditions_dict, max_conditions_lengths,
                  automatic_batching):
@@ -64,45 +69,68 @@ class PinaTensorDataset(PinaDataset):
                 in v.keys()} for k, v in self.conditions_dict.items()
         }
 
-    def _getitem_list(self, idx):
+    def fetch_from_idx_list(self, idx):
         to_return_dict = {}
         for condition, data in self.conditions_dict.items():
             cond_idx = idx[:self.max_conditions_lengths[condition]]
             condition_len = self.conditions_length[condition]
             if self.length > condition_len:
-                cond_idx = [idx%condition_len for idx in cond_idx]
+                cond_idx = [idx % condition_len for idx in cond_idx]
             to_return_dict[condition] = {k: v[cond_idx]
                                          for k, v in data.items()}
         return to_return_dict
 
+    @staticmethod
+    def _getitem_list(idx):
+        return idx
+
     def get_all_data(self):
         index = [i for i in range(len(self))]
-        return self._getitem_list(index)
+        return self.fetch_from_idx_list(index)
 
     def __getitem__(self, idx):
         return self._getitem_func(idx)
 
+
 class PinaGraphDataset(PinaDataset):
     pass
-    """
-    def __init__(self, conditions_dict, max_conditions_lengths):
+'''
+    def __init__(self, conditions_dict, max_conditions_lengths,
+                 automatic_batching):
         super().__init__(conditions_dict, max_conditions_lengths)
+        if automatic_batching:
+            self._getitem_func = self._getitem_int
+        else:
+            self._getitem_func = self._getitem_list
 
-    def __getitem__(self, idx):
-        
-        Getitem method for large batch size
-        
+    def fetch_from_idx_list(self, idx):
         to_return_dict = {}
         for condition, data in self.conditions_dict.items():
             cond_idx = idx[:self.max_conditions_lengths[condition]]
             condition_len = self.conditions_length[condition]
             if self.length > condition_len:
-                cond_idx = [idx%condition_len for idx in cond_idx]
+                cond_idx = [idx % condition_len for idx in cond_idx]
             to_return_dict[condition] = {k: Batch.from_data_list([v[i]
-                                            for i in cond_idx])
-                            if isinstance(v, list)
-                            else v[cond_idx].tensor.reshape(-1, v.size(-1))
-                        for k, v in data.items()
-                    }
+                                                                  for i in cond_idx])
+            if isinstance(v, list)
+            else v[cond_idx]
+                                         for k, v in data.items()
+                                         }
         return to_return_dict
-    """
+
+    def _getitem_list(self, idx):
+        return idx
+
+    def _getitem_int(self, idx):
+        return {
+            k: {k_data: v[k_data][idx % len(v['input_points'])] for k_data
+                in v.keys()} for k, v in self.conditions_dict.items()
+        }
+
+    def get_all_data(self):
+        index = [i for i in range(len(self))]
+        return self.fetch_from_idx_list(index)
+
+    def __getitem__(self, idx):
+        return self._getitem_func(idx)
+'''