Add TorchScript support for Node2Vec (#6726)

The original jit test of `Node2Vec` model is not really using
TorchScript, just fix this.

---------

Co-authored-by: Matthias Fey <matthias.fey@tu-dortmund.de>

CHANGELOG.md

@@ -7,6 +7,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
+- Added TorchScript support to the `Node2Vec` model ([#6726](https://github.com/pyg-team/pytorch_geometric/pull/6726))
 - Added `utils.to_edge_index` to convert sparse tensors to edge indices and edge attributes ([#6728](https://github.com/pyg-team/pytorch_geometric/issues/6728))
 - Fixed expected data format in `PolBlogs` dataset ([#6714](https://github.com/pyg-team/pytorch_geometric/issues/6714))
 - Added `SimpleConv` to perform non-trainable propagation ([#6718](https://github.com/pyg-team/pytorch_geometric/pull/6718))

test/nn/models/test_node2vec.py

@@ -22,7 +22,7 @@ def test_node2vec():
     assert 0 <= acc and acc <= 1
 
     if is_full_test():
-        jit = torch.jit.export(model)
+        jit = torch.jit.script(model)
 
         assert jit(torch.arange(3)).size() == (3, 16)
 

torch_geometric/nn/models/node2vec.py

@@ -14,8 +14,6 @@
 except ImportError:
     random_walk = None
 
-EPS = 1e-15
-
 
 class Node2Vec(torch.nn.Module):
     r"""The Node2Vec model from the
@@ -71,7 +69,7 @@ def __init__(
         row, col = edge_index
         self.adj = SparseTensor(row=row, col=col, sparse_sizes=(N, N))
         self.adj = self.adj.to('cpu')
-
+        self.EPS = 1e-15
         assert walk_length >= context_size
 
         self.embedding_dim = embedding_dim
@@ -99,6 +97,7 @@ def loader(self, **kwargs) -> DataLoader:
         return DataLoader(range(self.adj.sparse_size(0)),
                           collate_fn=self.sample, **kwargs)
 
+    @torch.jit.export
     def pos_sample(self, batch: Tensor) -> Tensor:
         batch = batch.repeat(self.walks_per_node)
         rowptr, col, _ = self.adj.csr()
@@ -112,6 +111,7 @@ def pos_sample(self, batch: Tensor) -> Tensor:
             walks.append(rw[:, j:j + self.context_size])
         return torch.cat(walks, dim=0)
 
+    @torch.jit.export
     def neg_sample(self, batch: Tensor) -> Tensor:
         batch = batch.repeat(self.walks_per_node * self.num_negative_samples)
 
@@ -125,11 +125,13 @@ def neg_sample(self, batch: Tensor) -> Tensor:
             walks.append(rw[:, j:j + self.context_size])
         return torch.cat(walks, dim=0)
 
+    @torch.jit.export
     def sample(self, batch: Tensor) -> Tuple[Tensor, Tensor]:
         if not isinstance(batch, Tensor):
             batch = torch.tensor(batch)
         return self.pos_sample(batch), self.neg_sample(batch)
 
+    @torch.jit.export
     def loss(self, pos_rw: Tensor, neg_rw: Tensor) -> Tensor:
         r"""Computes the loss given positive and negative random walks."""
 
@@ -142,7 +144,7 @@ def loss(self, pos_rw: Tensor, neg_rw: Tensor) -> Tensor:
                                                     self.embedding_dim)
 
         out = (h_start * h_rest).sum(dim=-1).view(-1)
-        pos_loss = -torch.log(torch.sigmoid(out) + EPS).mean()
+        pos_loss = -torch.log(torch.sigmoid(out) + self.EPS).mean()
 
         # Negative loss.
         start, rest = neg_rw[:, 0], neg_rw[:, 1:].contiguous()
@@ -153,7 +155,7 @@ def loss(self, pos_rw: Tensor, neg_rw: Tensor) -> Tensor:
                                                     self.embedding_dim)
 
         out = (h_start * h_rest).sum(dim=-1).view(-1)
-        neg_loss = -torch.log(1 - torch.sigmoid(out) + EPS).mean()
+        neg_loss = -torch.log(1 - torch.sigmoid(out) + self.EPS).mean()
 
         return pos_loss + neg_loss