lot of changes, no cleanup on setup, new node2vec class structure, need to setup sampler

Author: thunderock

Makefile

@@ -9,7 +9,7 @@ clean:
 
 
 .PHONY: setup
-setup: clean
+setup:
 	@echo "setting up..."
 ifeq ($(OS),Darwin)
 	@echo "Mac"
@@ -20,7 +20,7 @@ endif
 	@poetry install --only main -vvv
 
 .PHONY: setup_all
-setup_all: clean
+setup_all:
 	@echo "setting up..."
 	@poetry config virtualenvs.in-project true
 	@poetry install -vvv

graph_ml/models/gensim_node2vec.py

@@ -1,3 +1,21 @@
 import numpy as np
 import gensim
 from ..utils import config, torch_utils
+
+from ..models.node2vec import Node2Vec
+
+class GensimNode2Vec(Node2Vec):
+    def __init__(self, **params):
+        super().__init__(**params)
+        self.model_params = {
+            "vector_size": self.embedding_dim,
+            "window": self.context_size,
+            "min_count": 0,
+            "sg": 1,
+            "hs": 0,
+            "negative": 1,
+            "ns_exponent": 0.5,
+            "epochs": 1,
+            "workers": self.num_workers
+        }
+

graph_ml/models/node2vec.py

@@ -9,41 +9,19 @@
 class Node2Vec(object):
     def __init__(self, adj_list, embedding_dim, walk_length, context_size, device=config.DEVICE,
                  logging=config.LOGGING, **params):
-        edge_index = torch_utils.adj_list_to_edge_index(adj_list)
-        self.model = PyGNode2Vec(
-            edge_index, embedding_dim, walk_length, context_size, **params
-        ).to(device)
+        self.adj_list = adj_list
         self.num_workers = config.WORKER_COUNT
         self.logging = logging
-        self.loader = self.optimizer = None
+        self.embedding_dim = embedding_dim
+        self.walk_length = walk_length
+        self.context_size = context_size
+        self.device = device
 
     def fit(self, epochs=1, learning_rate=.1, batch_size=128):
-
-        # TODO (ashutosh): check if training two times works
-        self.loader = self.model.loader(
-            batch_size=batch_size, shuffle=True, num_workers=self.num_workers
-        )
-        self.optimizer = torch.optim.SparseAdam(self.model.parameters(), lr=learning_rate)
-        self.model.train()
-        total_loss = [0] * epochs
-        for epoch in range(epochs):
-            for pos_rw, neg_rw in self.loader:
-                self.optimizer.zero_grad()
-                loss = self.model.loss(pos_rw.to(self.model.device), neg_rw.to(self.model.device))
-                loss.backward()
-                self.optimizer.step()
-                total_loss[epoch] += loss.item()
-            total_loss[epoch] /= len(self.loader)
-            if self.logging:
-                print(f"Epoch: {epoch}, Loss: {total_loss[epoch]}")
-        return sum(total_loss) / epochs
+        return self._fit(epochs, learning_rate, batch_size)
 
     def transform(self, nodes=None, type_=np.ndarray):
-        if nodes is None:
-            nodes = torch.arange(self.model.num_nodes)
-        if type_ is np.ndarray:
-            return self.model(nodes).detach().cpu().numpy()
-        return self.model(nodes).detach()
+        return self._transform(nodes, type_)
 
     def fit_transform(self, epochs=1, learning_rate=.1, batch_size=128, nodes=None, type_=np.ndarray):
         self.fit(epochs, learning_rate, batch_size)

graph_ml/models/torch_node2vec.py

@@ -0,0 +1,56 @@
+import numpy as np
+from torch_geometric.nn import Node2Vec as PyGNode2Vec
+import torch
+from ..models.node2vec import Node2Vec
+from ..utils import config, torch_utils
+
+
+class TorchNode2Vec(Node2Vec):
+    def __init__(self, **params):
+        super().__init__(**params)
+        self.model = PyGNode2Vec(
+            edge_index=self.edge_index,
+            embedding_dim=self.embedding_dim,
+            walk_length=self.walk_length,
+            context_size=self.context_size,
+            sparse=True,
+            **params
+        ).to(self.device)
+        self.loader = self.optimizer = None
+
+    @property
+    def edge_index(self):
+        # should not be called too often, no caching here
+        return torch_utils.adj_list_to_edge_index(self.adj_list)
+
+
+    def _fit(self, epochs, learning_rate, batch_size, shuffle=True):
+
+        # TODO (ashutosh): check if training two times works
+        self.loader = self.model.loader(
+            batch_size=batch_size, shuffle=shuffle, num_workers=self.num_workers
+        )
+        self.optimizer = torch.optim.SparseAdam(self.model.parameters(), lr=learning_rate)
+        self.model.train()
+        total_loss = [0] * epochs
+        for epoch in range(epochs):
+            for pos_rw, neg_rw in self.loader:
+                self.optimizer.zero_grad()
+                loss = self.model.loss(pos_rw.to(self.model.device), neg_rw.to(self.device))
+                loss.backward()
+                self.optimizer.step()
+                total_loss[epoch] += loss.item()
+            total_loss[epoch] /= len(self.loader)
+            if self.logging:
+                print(f"Epoch: {epoch}, Loss: {total_loss[epoch]}")
+        return self
+
+    def _transform(self, nodes=None, type_=np.ndarray):
+        self.model.eval()
+        if nodes is None:
+            nodes = torch.arange(self.num_nodes, device=self.device)
+        with torch.no_grad():
+            emb = self.model(torch.tensor(nodes, device=self.device)).detach()
+        if type_ is np.ndarray:
+            return emb.cpu().numpy()
+        return emb

graph_ml/transformations/samplers/random_walk_sampler.py

@@ -18,6 +18,17 @@
 
 GPU_AVAILABLE = DEVICE_TYPE in ["cuda", "mps"]
 
+try:
+    import pyg_lib # noqa
+    WITH_PYG_LIB = True
+except:
+    WITH_PYG_LIB = False
+
+try:
+    import torch_cluster # noqa
+    WITH_TORCH_CLUSTER = True
+except:
+    WITH_TORCH_CLUSTER = False
 
 def get_formatted_os():
     if PLATFORM == "linux":