Repository: pyg-team/pytorch_geometric
Pull request #8546 (2023-12-05 22:57:41)

Example of a recommender system
Hi Everyone,

I'm adding a recommender system example with the following salient features

1. Dataset MovieLens – a heterogenous use case
2. Demonstrates the use of edge based temporal sampling
3. Visualization
  t-SNE based visualization (--visualize_emb)
4. Uses torch_geometric.nn.pool.MIPSKNNIndex for getting recommedations
5. Integration of the LinkPred metrics -- precision@k and ndcg@k

Thanks,
Poovaiah


codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8546?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Comparison is base [(`dd1a5be`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/dd1a5be993e902593a10767fb124168a4bc9c5b5?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) 88.52% compared to head [(`066dc0c`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8546?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) 88.52%.

> :exclamation: Current head 066dc0c differs from pull request most recent head a61b6b2. Consider uploading reports for the commit a61b6b2 to get more accurate results


<details><summary>Additional details and impacted files</summary>


```diff
@@           Coverage Diff           @@
##           master    #8546   +/-   ##
=======================================
  Coverage   88.52%   88.52%           
=======================================
  Files         479      479           
  Lines       29755    29755           
=======================================
  Hits        26340    26340           
  Misses       3415     3415           
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8546?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team).


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 694c00dfa..454028d5a 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
+- Added an example for recommender systems, including k-NN search and retrieval metrics ([#8546](https://github.com/pyg-team/pytorch_geometric/pull/8546))
 - Added multi-GPU evaluation in distributed sampling example ([#8880](https://github.com/pyg-team/pytorch_geometric/pull/8880))
 - Added end-to-end example for distributed CPU training ([#8713](https://github.com/pyg-team/pytorch_geometric/pull/8713))
 - Added PyTorch 2.2 support ([#8857](https://github.com/pyg-team/pyg-lib/pull/8857))
diff --git a/examples/hetero/recommender_system.py b/examples/hetero/recommender_system.py
new file mode 100644
index 000000000..672ef9029
--- /dev/null
+++ b/examples/hetero/recommender_system.py
@@ -0,0 +1,241 @@
+import argparse
+import os.path as osp
+
+import torch
+import torch.nn.functional as F
+from tqdm import tqdm
+
+import torch_geometric.transforms as T
+from torch_geometric import EdgeIndex
+from torch_geometric.datasets import MovieLens
+from torch_geometric.loader import LinkNeighborLoader, NeighborLoader
+from torch_geometric.nn import MIPSKNNIndex, SAGEConv, to_hetero
+from torch_geometric.nn.metrics import (
+    LinkPredMAP,
+    LinkPredPrecision,
+    LinkPredRecall,
+)
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--k', type=int, default=20, help='Number of predictions')
+args = parser.parse_args()
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+path = osp.join(osp.dirname(osp.realpath(__file__)), '../../data/MovieLens')
+data = MovieLens(path, model_name='all-MiniLM-L6-v2')[0]
+
+# Add user node features for message passing:
+data['user'].x = torch.eye(data['user'].num_nodes)
+del data['user'].num_nodes
+
+# Only use edges with high ratings (>= 4):
+mask = data['user', 'rates', 'movie'].edge_label >= 4
+data['user', 'movie'].edge_index = data['user', 'movie'].edge_index[:, mask]
+data['user', 'movie'].time = data['user', 'movie'].time[mask]
+del data['user', 'movie'].edge_label  # Drop rating information from graph.
+
+# Add a reverse ('movie', 'rev_rates', 'user') relation for message passing:
+data = T.ToUndirected()(data)
+
+# Perform a temporal link-level split into training, validation and test edges:
+edge_label_index = data['user', 'movie'].edge_index
+time = data['user', 'movie'].time
+
+perm = time.argsort()
+train_index = perm[:int(0.8 * perm.numel())]
+val_index = perm[int(0.8 * perm.numel()):int(0.9 * perm.numel())]
+test_index = perm[int(0.9 * perm.numel()):]
+
+kwargs = dict(  # Shared data loader arguments:
+    data=data,
+    num_neighbors=[5, 5, 5],
+    batch_size=256,
+    time_attr='time',
+    num_workers=4,
+    persistent_workers=True,
+    temporal_strategy='last',
+)
+
+train_loader = LinkNeighborLoader(
+    edge_label_index=(('user', 'movie'), edge_label_index[:, train_index]),
+    edge_label_time=time[train_index] - 1,  # No leakage.
+    neg_sampling=dict(mode='binary', amount=2),
+    shuffle=True,
+    **kwargs,
+)
+
+# During validation and testing, we sample node-level subgraphs from both
+# endpoints to retrieve their embeddings.
+# This allows us to do efficient k-NN search on top of embeddings:
+src_loader = NeighborLoader(
+    input_nodes='user',
+    input_time=(time[val_index].min() - 1).repeat(data['user'].num_nodes),
+    **kwargs,
+)
+dst_loader = NeighborLoader(
+    input_nodes='movie',
+    input_time=(time[val_index].min() - 1).repeat(data['movie'].num_nodes),
+    **kwargs,
+)
+
+# Save validation/test edges and the edges we want to exclude when evaluating:
+sparse_size = (data['user'].num_nodes, data['movie'].num_nodes)
+
+val_edge_label_index = EdgeIndex(
+    edge_label_index[:, val_index].to(device),
+    sparse_size=sparse_size,
+).sort_by('row')[0]
+val_exclude_links = EdgeIndex(
+    edge_label_index[:, train_index].to(device),
+    sparse_size=sparse_size,
+).sort_by('row')[0]
+
+test_edge_label_index = EdgeIndex(
+    edge_label_index[:, test_index].to(device),
+    sparse_size=sparse_size,
+).sort_by('row')[0]
+test_exclude_links = EdgeIndex(
+    edge_label_index[:, torch.cat([train_index, val_index])].to(device),
+    sparse_size=sparse_size,
+).sort_by('row')[0]
+
+
+class GNN(torch.nn.Module):
+    def __init__(self, hidden_channels):
+        super().__init__()
+        self.conv1 = SAGEConv((-1, -1), hidden_channels)
+        self.conv2 = SAGEConv((-1, -1), hidden_channels)
+        self.conv3 = SAGEConv((-1, -1), hidden_channels)
+
+    def forward(self, x, edge_index):
+        x = self.conv1(x, edge_index).relu()
+        x = self.conv2(x, edge_index).relu()
+        x = self.conv3(x, edge_index)
+        return x
+
+
+class InnerProductDecoder(torch.nn.Module):
+    def forward(self, x_dict, edge_label_index):
+        x_src = x_dict['user'][edge_label_index[0]]
+        x_dst = x_dict['movie'][edge_label_index[1]]
+        return (x_src * x_dst).sum(dim=-1)
+
+
+class Model(torch.nn.Module):
+    def __init__(self, hidden_channels):
+        super().__init__()
+        self.encoder = GNN(hidden_channels)
+        self.encoder = to_hetero(self.encoder, data.metadata(), aggr='sum')
+        self.decoder = InnerProductDecoder()
+
+    def forward(self, x_dict, edge_index_dict, edge_label_index):
+        x_dict = self.encoder(x_dict, edge_index_dict)
+        return self.decoder(x_dict, edge_label_index)
+
+
+model = Model(hidden_channels=64).to(device)
+optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
+
+
+def train():
+    model.train()
+
+    total_loss = total_examples = 0
+    for batch in tqdm(train_loader):
+        batch = batch.to(device)
+        optimizer.zero_grad()
+
+        out = model(
+            batch.x_dict,
+            batch.edge_index_dict,
+            batch['user', 'movie'].edge_label_index,
+        )
+        y = batch['user', 'movie'].edge_label
+
+        loss = F.binary_cross_entropy_with_logits(out, y)
+        loss.backward()
+        optimizer.step()
+
+        total_loss += float(loss) * y.numel()
+        total_examples += y.numel()
+
+    return total_loss / total_examples
+
+
+@torch.no_grad()
+def test(edge_label_index, exclude_links):
+    model.eval()
+
+    dst_embs = []
+    for batch in dst_loader:  # Collect destination node/movie embeddings:
+        batch = batch.to(device)
+        emb = model.encoder(batch.x_dict, batch.edge_index_dict)['movie']
+        emb = emb[:batch['movie'].batch_size]
+        dst_embs.append(emb)
+    dst_emb = torch.cat(dst_embs, dim=0)
+    del dst_embs
+
+    # Instantiate k-NN index based on maximum inner product search (MIPS):
+    mips = MIPSKNNIndex(dst_emb)
+
+    # Initialize metrics:
+    map_metric = LinkPredMAP(k=args.k).to(device)
+    precision_metric = LinkPredPrecision(k=args.k).to(device)
+    recall_metric = LinkPredRecall(k=args.k).to(device)
+
+    num_processed = 0
+    for batch in src_loader:  # Collect source node/user embeddings:
+        batch = batch.to(device)
+
+        # Compute user embeddings:
+        emb = model.encoder(batch.x_dict, batch.edge_index_dict)['user']
+        emb = emb[:batch['user'].batch_size]
+
+        # Filter labels/exclusion by current batch:
+        _edge_label_index = edge_label_index.sparse_narrow(
+            dim=0,
+            start=num_processed,
+            length=emb.size(0),
+        )
+        _exclude_links = exclude_links.sparse_narrow(
+            dim=0,
+            start=num_processed,
+            length=emb.size(0),
+        )
+        num_processed += emb.size(0)
+
+        # Perform MIPS search:
+        _, pred_index_mat = mips.search(emb, args.k, _exclude_links)
+
+        # Update retrieval metrics:
+        map_metric.update(pred_index_mat, _edge_label_index)
+        precision_metric.update(pred_index_mat, _edge_label_index)
+        recall_metric.update(pred_index_mat, _edge_label_index)
+
+    return (
+        float(map_metric.compute()),
+        float(precision_metric.compute()),
+        float(recall_metric.compute()),
+    )
+
+
+for epoch in range(1, 21):
+    train_loss = train()
+    print(f'Epoch: {epoch:02d}, Loss: {train_loss:.4f}')
+    val_map, val_precision, val_recall = test(
+        val_edge_label_index,
+        val_exclude_links,
+    )
+    print(f'Val MAP@{args.k}: {val_map:.4f}, '
+          f'Val Precision@{args.k}: {val_precision:.4f}, '
+          f'Val Recall@{args.k}: {val_recall:.4f}')
+
+print('Finished training! Evaluating on test set:')
+test_map, test_precision, test_recall = test(
+    test_edge_label_index,
+    test_exclude_links,
+)
+print(f'Test MAP@{args.k}: {test_map:.4f}, '
+      f'Test Precision@{args.k}: {test_precision:.4f}, '
+      f'Test Recall@{args.k}: {test_recall:.4f}')
diff --git a/torch_geometric/edge_index.py b/torch_geometric/edge_index.py
index 8ab684da3..cd11748ae 100644
--- a/torch_geometric/edge_index.py
+++ b/torch_geometric/edge_index.py
@@ -947,6 +947,93 @@ class EdgeIndex(Tensor):
         """
         return matmul(self, other, input_value, other_value, reduce, transpose)
 
+    def sparse_narrow(
+        self,
+        dim: int,
+        start: Union[int, Tensor],
+        length: int,
+    ) -> 'EdgeIndex':
+        r"""Returns a new :class:`EdgeIndex` that is a narrowed version of
+        itself. Narrowing is performed by interpreting :class:`EdgeIndex` as a
+        sparse matrix of shape :obj:`(num_rows, num_cols)`.
+
+        In contrast to :meth:`torch.narrow`, the returned tensor does not share
+        the same underlying storage anymore.
+
+        Args:
+            dim (int): The dimension along which to narrow.
+            start (int or torch.Tensor): Index of the element to start the
+                narrowed dimension from.
+            length (int): Length of the narrowed dimension.
+        """
+        dim = dim + 2 if dim < 0 else dim
+        if dim != 0 and dim != 1:
+            raise ValueError(f"Expected dimension to be 0 or 1 (got {dim})")
+
+        if start < 0:
+            raise ValueError(f"Expected 'start' value to be positive "
+                             f"(got {start})")
+
+        if dim == 0:
+            (rowptr, col), _ = self.get_csr()
+            rowptr = rowptr.narrow(0, start, length + 1)
+
+            if rowptr.numel() < 2:
+                row, col = self[0, :0], self[1, :0]
+                rowptr = None
+                num_rows = 0
+            else:
+                col = col[rowptr[0]:rowptr[-1]]
+                rowptr = rowptr - rowptr[0]
+                num_rows = rowptr.numel() - 1
+
+                row = torch.arange(
+                    num_rows,
+                    dtype=col.dtype,
+                    device=col.device,
+                ).repeat_interleave(
+                    rowptr.diff(),
+                    output_size=col.numel(),
+                )
+
+            edge_index = EdgeIndex(
+                torch.stack([row, col], dim=0),
+                sparse_size=(num_rows, self.sparse_size(1)),
+                sort_order='row',
+            )
+            edge_index._indptr = rowptr
+            return edge_index
+
+        else:  # dim == 0:
+            (colptr, row), _ = self.get_csc()
+            colptr = colptr.narrow(0, start, length + 1)
+
+            if colptr.numel() < 2:
+                row, col = self[0, :0], self[1, :0]
+                colptr = None
+                num_cols = 0
+            else:
+                row = row[colptr[0]:colptr[-1]]
+                colptr = colptr - colptr[0]
+                num_cols = colptr.numel() - 1
+
+                col = torch.arange(
+                    num_cols,
+                    dtype=row.dtype,
+                    device=row.device,
+                ).repeat_interleave(
+                    colptr.diff(),
+                    output_size=row.numel(),
+                )
+
+            edge_index = EdgeIndex(
+                torch.stack([row, col], dim=0),
+                sparse_size=(self.sparse_size(0), num_cols),
+                sort_order='col',
+            )
+            edge_index._indptr = colptr
+            return edge_index
+
     @classmethod
     def __torch_function__(
         cls: Type,


Test patch:
diff --git a/test/test_edge_index.py b/test/test_edge_index.py
index 3e9437108..bd168c6d2 100644
--- a/test/test_edge_index.py
+++ b/test/test_edge_index.py
@@ -989,6 +989,40 @@ def test_matmul(without_extensions, device):
     assert torch.allclose(out, expected)
 
 
+@withCUDA
+def test_sparse_narrow(device):
+    adj = EdgeIndex(
+        [[0, 1, 1, 2], [1, 0, 2, 1]],
+        device=device,
+        sort_order='row',
+    )
+
+    out = adj.sparse_narrow(dim=0, start=1, length=1)
+    assert out.equal(torch.tensor([[0, 0], [0, 2]], device=device))
+    assert out.sparse_size() == (1, None)
+    assert out.sort_order == 'row'
+    assert out._indptr.equal(torch.tensor([0, 2], device=device))
+
+    out = adj.sparse_narrow(dim=0, start=2, length=0)
+    assert out.equal(torch.tensor([[], []], device=device))
+    assert out.sparse_size() == (0, None)
+    assert out.sort_order == 'row'
+    assert out._indptr is None
+
+    out = adj.sparse_narrow(dim=1, start=1, length=1)
+    assert (out.equal(torch.tensor([[0, 2], [0, 0]], device=device))
+            or out.equal(torch.tensor([[2, 0], [0, 0]], device=device)))
+    assert out.sparse_size() == (3, 1)
+    assert out.sort_order == 'col'
+    assert out._indptr.equal(torch.tensor([0, 2], device=device))
+
+    out = adj.sparse_narrow(dim=1, start=2, length=0)
+    assert out.equal(torch.tensor([[], []], device=device))
+    assert out.sparse_size() == (3, 0)
+    assert out.sort_order == 'col'
+    assert out._indptr is None
+
+
 @withCUDA
 @pytest.mark.parametrize('dtype', DTYPES)
 def test_save_and_load(dtype, device, tmp_path):