feat: graph saint sampler

Author: tien02

biomedkg/data_module.py

@@ -1,6 +1,6 @@
 import torch_geometric.transforms as T
 from lightning import LightningDataModule
-from torch_geometric.loader import NeighborLoader
+from torch_geometric.loader import GraphSAINTRandomWalkSampler, NeighborLoader
 
 from biomedkg.data import dataset, node
 
@@ -86,22 +86,44 @@ def all_dataloader(self):
             num_workers=0,
         )
 
-    def train_dataloader(self):
-        return NeighborLoader(
-            data=self.train_data,
-            batch_size=self.batch_size,
-            num_neighbors=[30] * 3,
-            num_workers=0,
-            shuffle=True,
-        )
-
-    def val_dataloader(self):
-        return NeighborLoader(
-            data=self.val_data,
-            batch_size=self.batch_size,
-            num_neighbors=[30] * 3,
-            num_workers=0,
-        )
+    def train_dataloader(self, loader_type: str = "neighbor"):
+        assert loader_type in ["neighbor", "saint"]
+
+        if loader_type == "neighbor":
+            return NeighborLoader(
+                data=self.train_data,
+                batch_size=self.batch_size,
+                num_neighbors=[30] * 3,
+                num_workers=0,
+                shuffle=True,
+            )
+        elif loader_type == "saint":
+            return GraphSAINTRandomWalkSampler(
+                data=self.train_data,
+                batch_size=self.batch_size,
+                walk_length=10,
+                num_steps=1000,
+                num_workers=0,
+            )
+
+    def val_dataloader(self, loader_type: str = "neighbor"):
+        assert loader_type in ["neighbor", "saint"]
+
+        if loader_type == "neighbor":
+            return NeighborLoader(
+                data=self.val_data,
+                batch_size=self.batch_size,
+                num_neighbors=[30] * 3,
+                num_workers=0,
+            )
+        elif loader_type == "saint":
+            return GraphSAINTRandomWalkSampler(
+                data=self.val_data,
+                batch_size=self.batch_size,
+                walk_length=10,
+                num_steps=100,
+                num_workers=0,
+            )
 
     def test_dataloader(self):
         return NeighborLoader(

train_dpi.py

@@ -94,8 +94,8 @@ def main(cfg: DictConfig):
 
     trainer.fit(
         model=model,
-        train_dataloaders=data_module.train_dataloader(),
-        val_dataloaders=data_module.val_dataloader(),
+        train_dataloaders=data_module.train_dataloader(loader_type="saint"),
+        val_dataloaders=data_module.val_dataloader(loader_type="saint"),
     )
 
     test_args = {

train_gcl.py

@@ -107,8 +107,8 @@ def main(cfg: DictConfig):
 
     trainer.fit(
         model=model,
-        train_dataloaders=data_module.train_dataloader(),
-        val_dataloaders=data_module.val_dataloader(),
+        train_dataloaders=data_module.train_dataloader(loader_type="neighbor"),
+        val_dataloaders=data_module.val_dataloader(loader_type="neighbor"),
     )
 
     test_args = {

train_kge.py

@@ -87,8 +87,8 @@ def main(cfg: DictConfig):
 
     trainer.fit(
         model=model,
-        train_dataloaders=data_module.train_dataloader(),
-        val_dataloaders=data_module.val_dataloader(),
+        train_dataloaders=data_module.train_dataloader(loader_type="saint"),
+        val_dataloaders=data_module.val_dataloader(loader_type="saint"),
     )
 
     test_args = {