fix scaling and dataloader len

TUM-DAML · Dec 12, 2022 · 093316b · 093316b
1 parent 146c361
commit 093316b
Show file tree

Hide file tree

Showing 8 changed files with 13 additions and 18 deletions.
diff --git a/dataloaders/BaseLoader.py b/dataloaders/BaseLoader.py
@@ -30,6 +30,10 @@ def __getitem__(self, *args, **kwargs):
     def __len__(self, *args, **kwargs):
         raise NotImplementedError
 
+    @property
+    def loader_len(self):
+        return len(self)
+
     def __collate__(self, *args, **kwargs):
         raise NotImplementedError
 

diff --git a/dataloaders/GraphSAINTRWSampler.py b/dataloaders/GraphSAINTRWSampler.py
@@ -196,6 +196,7 @@ def __filename__(self):
     def __len__(self):
         return len(self.output_indices)
 
+    @property
     def loader_len(self):
         return ceil(len(self.output_indices) / self.batch_size)
 

diff --git a/dataloaders/IBMBRandLoader.py b/dataloaders/IBMBRandLoader.py
@@ -65,6 +65,7 @@ def __getitem__(self, idx):
     def __len__(self):
         return len(self.out_aux_pairs)
 
+    @property
     def loader_len(self):
         return ceil(len(self.out_aux_pairs) / self.batch_size)
 

diff --git a/dataloaders/LADIESSampler.py b/dataloaders/LADIESSampler.py
@@ -67,6 +67,7 @@ def __getitem__(self, idx):
     def __len__(self):
         return len(self.output_indices)
 
+    @property
     def loader_len(self):
         return ceil(len(self.output_indices) / self.batch_size)
 

diff --git a/dataloaders/NeighborSamplingLoader.py b/dataloaders/NeighborSamplingLoader.py
@@ -35,6 +35,7 @@ def __getitem__(self, idx):
     def __len__(self):
         return len(self.node_idx)
 
+    @property
     def loader_len(self):
         return ceil(len(self.node_idx) / self.batch_size)
 

diff --git a/dataloaders/ShaDowLoader.py b/dataloaders/ShaDowLoader.py
@@ -65,6 +65,7 @@ def __getitem__(self, idx):
     def __len__(self):
         return len(self.out_aux_pairs)
 
+    @property
     def loader_len(self):
         return ceil(len(self.out_aux_pairs) / self.batch_size)
 

diff --git a/train/prefetch_generators.py b/train/prefetch_generators.py
@@ -1,12 +1,6 @@
 import queue
 import threading
 
-from dataloaders.GraphSAINTRWSampler import SaintRWValSampler
-from dataloaders.IBMBRandLoader import IBMBRandLoader
-from dataloaders.ShaDowLoader import ShaDowLoader
-from dataloaders.LADIESSampler import LADIESSampler
-from dataloaders.NeighborSamplingLoader import NeighborSamplingLoader
-
 
 class BaseGenerator(threading.Thread):
     def __init__(self, max_prefetch=1, device='cuda'):
@@ -37,13 +31,6 @@ def __init__(self, dataloader, max_prefetch=1, device='cuda'):
 
     def run(self):
         for i, graph in enumerate(self.dataloader):
-            if isinstance(self.dataloader, (SaintRWValSampler,
-                                            ShaDowLoader,
-                                            IBMBRandLoader,
-                                            LADIESSampler,
-                                            NeighborSamplingLoader)):
-                stop_signal = i == self.dataloader.loader_len() - 1
-            else:
-                stop_signal = i == len(self.dataloader) - 1
+            stop_signal = i == self.dataloader.loader_len - 1
             self.queue.put((graph.to(self.device, non_blocking=True), stop_signal))
         self.queue.put(None)
diff --git a/train/trainer.py b/train/trainer.py
@@ -44,8 +44,7 @@ def __init__(self,
 
     def get_loss_scaling(self, len_loader: int):
         micro_batch = int(min(self.micro_batch, len_loader))
-        num_batches = ceil(len_loader / self.batch_size)
-        loss_scaling_lst = [micro_batch] * (num_batches // micro_batch) + [num_batches % micro_batch]
+        loss_scaling_lst = [micro_batch] * (len_loader // micro_batch) + [len_loader % micro_batch]
         return loss_scaling_lst, micro_batch
 
     def train(self,
@@ -90,7 +89,7 @@ def train(self,
 
             # train
             model.train()
-            loss_scaling_lst, cur_micro_batch = self.get_loss_scaling(len(train_loader))
+            loss_scaling_lst, cur_micro_batch = self.get_loss_scaling(train_loader.loader_len)
             loader, next_loader = next_loader, None
 
             start_time = time.time()
@@ -385,7 +384,7 @@ def full_graph_inference(self,
         adj = BaseLoader.normalize_adjmat(adj, normalization='sym')
 
         outputs = model.chunked_pass(MyGraph(x=graph.x, adj=adj, idx=torch.from_numpy(mask)),
-                                     self.num_batches // self.batch_size).detach().numpy()
+                                     self.num_batches // self.batch_size).detach().numpy()  # an estimate of #chunks
 
         for cat in ['val', 'test']:
             nodes = val_nodes if cat == 'val' else test_nodes