Repository: pyg-team/pytorch_geometric
Pull request #9941 (2025-01-14 05:19:13)

Add `LinkPredMetricCollection`


Patch:
diff --git a/.github/labeler.yml b/.github/labeler.yml
index ae2f57cc9..8a2f60bbd 100644
--- a/.github/labeler.yml
+++ b/.github/labeler.yml
@@ -34,6 +34,10 @@ transform:
   - changed-files:
       - any-glob-to-any-file: "torch_geometric/transforms/**/*"
 
+metrics:
+  - changed-files:
+      - any-glob-to-any-file: "torch_geometric/metrics/**/*"
+
 utils:
   - changed-files:
       - any-glob-to-any-file: "torch_geometric/utils/**/*"
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 8691f4af5..50ff227f4 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,8 +7,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
-- Update Dockerfile to use latest from NVIDIA ([#9794](https://github.com/pyg-team/pytorch_geometric/pull/9794))
-- Added various GRetriever Architecture Benchmarking examples ([#9666](https://github.com/pyg-team/pytorch_geometric/pull/9666))
+- Added `LinkPredMetricCollection` ([#9941](https://github.com/pyg-team/pytorch_geometric/pull/9941))
+- Added various `GRetriever` architecture benchmarking examples ([#9666](https://github.com/pyg-team/pytorch_geometric/pull/9666))
 - Added `profiler.nvtxit` with some examples ([#9666](https://github.com/pyg-team/pytorch_geometric/pull/9666))
 - Added `loader.RagQueryLoader` with Remote Backend Example ([#9666](https://github.com/pyg-team/pytorch_geometric/pull/9666))
 - Added `data.LargeGraphIndexer` ([#9666](https://github.com/pyg-team/pytorch_geometric/pull/9666))
@@ -26,6 +26,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Changed
 
+- Updated Dockerfile to use latest from NVIDIA ([#9794](https://github.com/pyg-team/pytorch_geometric/pull/9794))
 - Dropped Python 3.8 support ([#9696](https://github.com/pyg-team/pytorch_geometric/pull/9606))
 - Added a check that confirms that custom edge types of `NumNeighbors` actually exist in the graph ([#9807](https://github.com/pyg-team/pytorch_geometric/pull/9807))
 
diff --git a/torch_geometric/metrics/__init__.py b/torch_geometric/metrics/__init__.py
index 1340829b8..e142a7dc1 100644
--- a/torch_geometric/metrics/__init__.py
+++ b/torch_geometric/metrics/__init__.py
@@ -1,6 +1,7 @@
 # flake8: noqa
 
 from .link_pred import (
+    LinkPredMetricCollection,
     LinkPredPrecision,
     LinkPredRecall,
     LinkPredF1,
@@ -10,6 +11,7 @@ from .link_pred import (
 )
 
 link_pred_metrics = [
+    'LinkPredMetricCollection',
     'LinkPredPrecision',
     'LinkPredRecall',
     'LinkPredF1',
diff --git a/torch_geometric/metrics/link_pred.py b/torch_geometric/metrics/link_pred.py
index a450bcd1a..fa5ff371a 100644
--- a/torch_geometric/metrics/link_pred.py
+++ b/torch_geometric/metrics/link_pred.py
@@ -1,4 +1,4 @@
-from typing import Optional, Tuple, Union
+from typing import Dict, List, Optional, Tuple, Union
 
 import torch
 from torch import Tensor
@@ -144,6 +144,103 @@ class LinkPredMetric(BaseMetric):
         return f'{self.__class__.__name__}(k={self.k})'
 
 
+class LinkPredMetricCollection(torch.nn.ModuleDict):
+    r"""A collection of metrics to reduce and speed-up computation of link
+    prediction metrics.
+
+    .. code-block:: python
+
+        from torch_geometric.metrics import (
+            LinkPredMAP,
+            LinkPredMetricCollection,
+            LinkPredPrecision,
+            LinkPredRecall,
+        )
+
+        metrics = LinkPredMetricCollection([
+            LinkPredMAP(k=10),
+            LinkPredPrecision(k=100),
+            LinkPredRecall(k=50),
+        ])
+
+        metrics.update(pred_index_mat, edge_label_index)
+        out = metrics.compute()
+        metrics.reset()
+
+        print(out)
+        >>> {'LinkPredMAP@10': tensor(0.375),
+        ...  'LinkPredPrecision@100': tensor(0.127),
+        ...  'LinkPredRecall@50': tensor(0.483)}
+
+    Args:
+        metrics: The link prediction metrics.
+    """
+    def __init__(
+        self,
+        metrics: Union[
+            List[LinkPredMetric],
+            Dict[str, LinkPredMetric],
+        ],
+    ) -> None:
+        super().__init__()
+
+        if isinstance(metrics, (list, tuple)):
+            metrics = {
+                f'{metric.__class__.__name__}@{metric.k}': metric
+                for metric in metrics
+            }
+        assert len(metrics) > 0
+        assert isinstance(metrics, dict)
+
+        for name, metric in metrics.items():
+            self[name] = metric
+
+    @property
+    def max_k(self) -> int:
+        r"""The maximum number of top-:math:`k` predictions to evaluate
+        against.
+        """
+        return max([metric.k for metric in self.values()])
+
+    def update(  # type: ignore
+        self,
+        pred_index_mat: Tensor,
+        edge_label_index: Union[Tensor, Tuple[Tensor, Tensor]],
+    ) -> None:
+        r"""Updates the state variables based on the current mini-batch
+        prediction.
+
+        :meth:`update` can be repeated multiple times to accumulate the results
+        of successive predictions, *e.g.*, inside a mini-batch training or
+        evaluation loop.
+
+        Args:
+            pred_index_mat (torch.Tensor): The top-:math:`k` predictions of
+                every example in the mini-batch with shape
+                :obj:`[batch_size, k]`.
+            edge_label_index (torch.Tensor): The ground-truth indices for every
+                example in the mini-batch, given in COO format of shape
+                :obj:`[2, num_ground_truth_indices]`.
+        """
+        pred_isin_mat, y_count = LinkPredMetric._prepare(
+            pred_index_mat, edge_label_index)
+        for metric in self.values():
+            metric._update_from_prepared(pred_isin_mat, y_count)
+
+    def compute(self) -> Dict[str, Tensor]:
+        r"""Computes the final metric values."""
+        return {name: metric.compute() for name, metric in self.items()}
+
+    def reset(self) -> None:
+        r"""Reset metric state variables to their default value."""
+        for metric in self.values():
+            metric.reset()
+
+    def __repr__(self) -> str:
+        names = [f'  {name}: {metric},\n' for name, metric in self.items()]
+        return f'{self.__class__.__name__}([\n{"".join(names)}])'
+
+
 class LinkPredPrecision(LinkPredMetric):
     r"""A link prediction metric to compute Precision @ :math:`k`.
 


Test patch:
diff --git a/test/metrics/test_link_pred_metric.py b/test/metrics/test_link_pred_metric.py
index 1dd30a352..6256eeb4d 100644
--- a/test/metrics/test_link_pred_metric.py
+++ b/test/metrics/test_link_pred_metric.py
@@ -6,6 +6,7 @@ import torch
 from torch_geometric.metrics import (
     LinkPredF1,
     LinkPredMAP,
+    LinkPredMetricCollection,
     LinkPredMRR,
     LinkPredNDCG,
     LinkPredPrecision,
@@ -25,7 +26,7 @@ def test_precision(num_src_nodes, num_dst_nodes, num_edges, batch_size, k):
 
     pred = torch.rand(num_src_nodes, num_dst_nodes)
     pred[row, col] += 0.3  # Offset positive links by a little.
-    top_k_pred_mat = pred.topk(k, dim=1)[1]
+    pred_index_mat = pred.topk(k, dim=1)[1]
 
     metric = LinkPredPrecision(k)
     assert str(metric) == f'LinkPredPrecision(k={k})'
@@ -39,7 +40,7 @@ def test_precision(num_src_nodes, num_dst_nodes, num_edges, batch_size, k):
         arange[node_id] = torch.arange(node_id.numel())
         y_batch = arange[y_batch]
 
-        metric.update(top_k_pred_mat[node_id], (y_batch, y_index))
+        metric.update(pred_index_mat[node_id], (y_batch, y_index))
 
     out = metric.compute()
     metric.reset()
@@ -48,96 +49,135 @@ def test_precision(num_src_nodes, num_dst_nodes, num_edges, batch_size, k):
     for i in range(num_src_nodes):  # Naive computation per node:
         y_index = col[row == i]
         if y_index.numel() > 0:
-            mask = torch.isin(top_k_pred_mat[i], y_index)
+            mask = torch.isin(pred_index_mat[i], y_index)
             precision = float(mask.sum() / k)
             values.append(precision)
     expected = torch.tensor(values).mean()
     assert torch.allclose(out, expected)
 
-    # Test with `k > pred_mat.size(1)`:
-    metric.update(top_k_pred_mat[:, :k - 1], edge_label_index)
+    # Test with `k > pred_index_mat.size(1)`:
+    metric.update(pred_index_mat[:, :k - 1], edge_label_index)
     metric.compute()
     metric.reset()
 
 
 def test_recall():
-    pred_mat = torch.tensor([[1, 0], [1, 2], [0, 2]])
+    pred_index_mat = torch.tensor([[1, 0], [1, 2], [0, 2]])
     edge_label_index = torch.tensor([[0, 0, 0, 2, 2], [0, 1, 2, 2, 1]])
 
     metric = LinkPredRecall(k=2)
     assert str(metric) == 'LinkPredRecall(k=2)'
-    metric.update(pred_mat, edge_label_index)
+    metric.update(pred_index_mat, edge_label_index)
     result = metric.compute()
 
     assert float(result) == pytest.approx(0.5 * (2 / 3 + 0.5))
 
-    # Test with `k > pred_mat.size(1)`:
-    metric.update(pred_mat[:, :1], edge_label_index)
+    # Test with `k > pred_index_mat.size(1)`:
+    metric.update(pred_index_mat[:, :1], edge_label_index)
     metric.compute()
     metric.reset()
 
 
 def test_f1():
-    pred_mat = torch.tensor([[1, 0], [1, 2], [0, 2]])
+    pred_index_mat = torch.tensor([[1, 0], [1, 2], [0, 2]])
     edge_label_index = torch.tensor([[0, 0, 0, 2, 2], [0, 1, 2, 2, 1]])
 
     metric = LinkPredF1(k=2)
     assert str(metric) == 'LinkPredF1(k=2)'
-    metric.update(pred_mat, edge_label_index)
+    metric.update(pred_index_mat, edge_label_index)
     result = metric.compute()
     assert float(result) == pytest.approx(0.6500)
 
-    # Test with `k > pred_mat.size(1)`:
-    metric.update(pred_mat[:, :1], edge_label_index)
+    # Test with `k > pred_index_mat.size(1)`:
+    metric.update(pred_index_mat[:, :1], edge_label_index)
     metric.compute()
     metric.reset()
 
 
 def test_map():
-    pred_mat = torch.tensor([[1, 0], [1, 2], [0, 2]])
+    pred_index_mat = torch.tensor([[1, 0], [1, 2], [0, 2]])
     edge_label_index = torch.tensor([[0, 0, 0, 2, 2], [0, 1, 2, 2, 1]])
 
     metric = LinkPredMAP(k=2)
     assert str(metric) == 'LinkPredMAP(k=2)'
-    metric.update(pred_mat, edge_label_index)
+    metric.update(pred_index_mat, edge_label_index)
     result = metric.compute()
     assert float(result) == pytest.approx(0.6250)
 
-    # Test with `k > pred_mat.size(1)`:
-    metric.update(pred_mat[:, :1], edge_label_index)
+    # Test with `k > pred_index_mat.size(1)`:
+    metric.update(pred_index_mat[:, :1], edge_label_index)
     metric.compute()
     metric.reset()
 
 
 def test_ndcg():
-    pred_mat = torch.tensor([[1, 0], [1, 2], [0, 2]])
+    pred_index_mat = torch.tensor([[1, 0], [1, 2], [0, 2]])
     edge_label_index = torch.tensor([[0, 0, 2, 2], [0, 1, 2, 1]])
 
     metric = LinkPredNDCG(k=2)
     assert str(metric) == 'LinkPredNDCG(k=2)'
-    metric.update(pred_mat, edge_label_index)
+    metric.update(pred_index_mat, edge_label_index)
     result = metric.compute()
 
     assert float(result) == pytest.approx(0.6934264)
 
-    # Test with `k > pred_mat.size(1)`:
-    metric.update(pred_mat[:, :1], edge_label_index)
+    # Test with `k > pred_index_mat.size(1)`:
+    metric.update(pred_index_mat[:, :1], edge_label_index)
     metric.compute()
     metric.reset()
 
 
 def test_mrr():
-    pred_mat = torch.tensor([[1, 0], [1, 2], [0, 2], [0, 1]])
+    pred_index_mat = torch.tensor([[1, 0], [1, 2], [0, 2], [0, 1]])
     edge_label_index = torch.tensor([[0, 0, 2, 2, 3], [0, 1, 2, 1, 2]])
 
     metric = LinkPredMRR(k=2)
     assert str(metric) == 'LinkPredMRR(k=2)'
-    metric.update(pred_mat, edge_label_index)
+    metric.update(pred_index_mat, edge_label_index)
     result = metric.compute()
 
     assert float(result) == pytest.approx((1 + 0.5 + 0) / 3)
 
-    # Test with `k > pred_mat.size(1)`:
-    metric.update(pred_mat[:, :1], edge_label_index)
+    # Test with `k > pred_index_mat.size(1)`:
+    metric.update(pred_index_mat[:, :1], edge_label_index)
     metric.compute()
     metric.reset()
+
+
+@pytest.mark.parametrize('num_src_nodes', [10])
+@pytest.mark.parametrize('num_dst_nodes', [50])
+@pytest.mark.parametrize('num_edges', [200])
+def test_link_pred_metric_collection(num_src_nodes, num_dst_nodes, num_edges):
+    metrics = [
+        LinkPredMAP(k=10),
+        LinkPredPrecision(k=100),
+        LinkPredRecall(k=50),
+    ]
+
+    row = torch.randint(0, num_src_nodes, (num_edges, ))
+    col = torch.randint(0, num_dst_nodes, (num_edges, ))
+    edge_label_index = torch.stack([row, col], dim=0)
+
+    pred = torch.rand(num_src_nodes, num_dst_nodes)
+    pred[row, col] += 0.3  # Offset positive links by a little.
+    pred_index_mat = pred.argsort(dim=1)
+
+    metric_collection = LinkPredMetricCollection(metrics)
+    assert str(metric_collection) == (
+        'LinkPredMetricCollection([\n'
+        '  LinkPredMAP@10: LinkPredMAP(k=10),\n'
+        '  LinkPredPrecision@100: LinkPredPrecision(k=100),\n'
+        '  LinkPredRecall@50: LinkPredRecall(k=50),\n'
+        '])')
+    assert metric_collection.max_k == 100
+
+    expected = {}
+    for metric in metrics:
+        metric.update(pred_index_mat[:, :metric.k], edge_label_index)
+        out = metric.compute()
+        expected[f'{metric.__class__.__name__}@{metric.k}'] = out
+        metric.reset()
+
+    metric_collection.update(pred_index_mat, edge_label_index)
+    assert metric_collection.compute() == expected
+    metric_collection.reset()