Repository: pyg-team/pytorch_geometric
Pull request #9632 (2024-08-29 14:03:11)

Mean reciprocal rank metric
Addresses #9631. Implements `LinkPredMRR` as a `LinkPredMetric`, with test included.


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index f58d082ed..023ae7cc1 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 the `LinkPredMRR` metric ([#9632](https://github.com/pyg-team/pytorch_geometric/pull/9632))
 - Added PyTorch 2.4 support ([#9594](https://github.com/pyg-team/pytorch_geometric/pull/9594))
 - Added `utils.normalize_edge_index` for symmetric/asymmetric normalization of graph edges ([#9554](https://github.com/pyg-team/pytorch_geometric/pull/9554))
 - Added the `RemoveSelfLoops` transformation ([#9562](https://github.com/pyg-team/pytorch_geometric/pull/9562))
diff --git a/README.md b/README.md
index 713c3acc7..17a51ff6d 100644
--- a/README.md
+++ b/README.md
@@ -383,9 +383,9 @@ where `${CUDA}` should be replaced by either `cpu`, `cu118`, `cu121`, or `cu124`
 
 |             | `cpu` | `cu118` | `cu121` | `cu124` |
 | ----------- | ----- | ------- | ------- | ------- |
-| **Linux**   | ✅     | ✅       | ✅       | ✅       |
-| **Windows** | ✅     | ✅       | ✅       | ✅       |
-| **macOS**   | ✅     |         |         |         |
+| **Linux**   | ✅    | ✅      | ✅      | ✅      |
+| **Windows** | ✅    | ✅      | ✅      | ✅      |
+| **macOS**   | ✅    |         |         |         |
 
 #### PyTorch 2.3
 
@@ -399,9 +399,9 @@ where `${CUDA}` should be replaced by either `cpu`, `cu118`, or `cu121` dependin
 
 |             | `cpu` | `cu118` | `cu121` |
 | ----------- | ----- | ------- | ------- |
-| **Linux**   | ✅     | ✅       | ✅       |
-| **Windows** | ✅     | ✅       | ✅       |
-| **macOS**   | ✅     |         |         |
+| **Linux**   | ✅    | ✅      | ✅      |
+| **Windows** | ✅    | ✅      | ✅      |
+| **macOS**   | ✅    |         |         |
 
 **Note:** Binaries of older versions are also provided for PyTorch 1.4.0, PyTorch 1.5.0, PyTorch 1.6.0, PyTorch 1.7.0/1.7.1, PyTorch 1.8.0/1.8.1, PyTorch 1.9.0, PyTorch 1.10.0/1.10.1/1.10.2, PyTorch 1.11.0, PyTorch 1.12.0/1.12.1, PyTorch 1.13.0/1.13.1, PyTorch 2.0.0/2.0.1, PyTorch 2.1.0/2.1.1/2.1.2, and PyTorch 2.2.0/2.2.1/2.2.2 (following the same procedure).
 **For older versions, you might need to explicitly specify the latest supported version number** or install via `pip install --no-index` in order to prevent a manual installation from source.
diff --git a/examples/multi_gpu/papers100m_gcn.py b/examples/multi_gpu/papers100m_gcn.py
index e7f6fe0dc..d7329e902 100644
--- a/examples/multi_gpu/papers100m_gcn.py
+++ b/examples/multi_gpu/papers100m_gcn.py
@@ -157,7 +157,17 @@ if __name__ == '__main__':
 
     args = parser.parse_args()
     wall_clock_start = time.perf_counter()
-
+    if args.n_devices == -1:
+        world_size = torch.cuda.device_count()
+    else:
+        world_size = args.n_devices
+    import psutil
+    gb_ram_needed = 190 + 200 * world_size
+    if (psutil.virtual_memory().total / (1024**3)) < gb_ram_needed:
+        print("Warning: may not have enough RAM to use this many GPUs.")
+        print("Consider upgrading RAM or using less GPUs if an error occurs.")
+        print("Estimated RAM Needed: ~" + str(gb_ram_needed))
+    print('Let\'s use', world_size, 'GPUs!')
     dataset = PygNodePropPredDataset(name='ogbn-papers100M',
                                      root='/datasets/ogb_datasets')
     split_idx = dataset.get_idx_split()
@@ -178,11 +188,6 @@ if __name__ == '__main__':
         )
 
     print("Data =", data)
-    if args.n_devices == -1:
-        world_size = torch.cuda.device_count()
-    else:
-        world_size = args.n_devices
-    print('Let\'s use', world_size, 'GPUs!')
     with tempfile.TemporaryDirectory() as tempdir:
         if world_size > 1:
             mp.spawn(
diff --git a/torch_geometric/metrics/__init__.py b/torch_geometric/metrics/__init__.py
index 95a198a25..1340829b8 100644
--- a/torch_geometric/metrics/__init__.py
+++ b/torch_geometric/metrics/__init__.py
@@ -1,7 +1,13 @@
 # flake8: noqa
 
-from .link_pred import (LinkPredPrecision, LinkPredRecall, LinkPredF1,
-                        LinkPredMAP, LinkPredNDCG)
+from .link_pred import (
+    LinkPredPrecision,
+    LinkPredRecall,
+    LinkPredF1,
+    LinkPredMAP,
+    LinkPredNDCG,
+    LinkPredMRR,
+)
 
 link_pred_metrics = [
     'LinkPredPrecision',
@@ -9,6 +15,7 @@ link_pred_metrics = [
     'LinkPredF1',
     'LinkPredMAP',
     'LinkPredNDCG',
+    'LinkPredMRR',
 ]
 
 __all__ = link_pred_metrics
diff --git a/torch_geometric/metrics/link_pred.py b/torch_geometric/metrics/link_pred.py
index 9ea644300..25967a171 100644
--- a/torch_geometric/metrics/link_pred.py
+++ b/torch_geometric/metrics/link_pred.py
@@ -216,3 +216,20 @@ class LinkPredNDCG(LinkPredMetric):
         out = dcg / idcg
         out[out.isnan() | out.isinf()] = 0.0
         return out
+
+
+class LinkPredMRR(LinkPredMetric):
+    r"""A link prediction metric to compute the MRR @ :math:`k` (Mean
+    Reciprocal Rank).
+
+    Args:
+        k (int): The number of top-:math:`k` predictions to evaluate against.
+    """
+    higher_is_better: bool = True
+
+    def _compute(self, pred_isin_mat: Tensor, y_count: Tensor) -> Tensor:
+        rank = pred_isin_mat.type(torch.uint8).argmax(dim=-1)
+        is_correct = pred_isin_mat.gather(1, rank.view(-1, 1)).view(-1)
+        reciprocals = 1.0 / (rank + 1)
+        reciprocals[~is_correct] = 0.0
+        return reciprocals


Test patch:
diff --git a/test/metrics/test_link_pred_metric.py b/test/metrics/test_link_pred_metric.py
index ee2f08a2b..f4bb5c1af 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,
+    LinkPredMRR,
     LinkPredNDCG,
     LinkPredPrecision,
     LinkPredRecall,
@@ -98,3 +99,15 @@ def test_ndcg():
     result = metric.compute()
 
     assert float(result) == pytest.approx(0.6934264)
+
+
+def test_mrr():
+    pred_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)
+    result = metric.compute()
+
+    assert float(result) == pytest.approx((1 + 0.5 + 0) / 3)