Repository: pyg-team/pytorch_geometric
Pull request #8298 (2023-10-31 13:31:53)

Add mean reciprocal rank (MRR) computation in `KGEModel`
Greetings, thank you for confirming this request.
I tried to add mean reciprocal rank (MRR) in the torch_geometric/nn/kge/base.py and examples/kge_fb15k_237.py.
(Please see the related discussion in [here](https://github.com/pyg-team/pytorch_geometric/discussions/8256))
If there are any problems, please forgive me and freely edit the code.
Thanks!

songsong0425: Are there any additional variables for the evaluation metrics? I couldn't understand why I faced this error after add `mrr` to base.py and kge_fb15k_237.py:
```
File "/home/songyeon/PyG_contribution/pytorch_geometric/examples/kge_fb15k_237.py", line 83, in <module>
    rank, mrr, hits = test(val_data)
ValueError: not enough values to unpack (expected 3, got 2)
```
rusty1s: > Are there any additional variables for the evaluation metrics? I couldn't understand why I faced this error after add mrr to base.py and kge_fb15k_237.py:

Looks like you didn't install the package with your corresponding changes?

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 709c91b20..6d1bf8db8 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
 
+- Support MRR computation in `KGEModel.test()` ([#8298](https://github.com/pyg-team/pytorch_geometric/pull/8298))
 - Added an example for model parallelism (`examples/multi_gpu/model_parallel.py`) ([#8309](https://github.com/pyg-team/pytorch_geometric/pull/8309))
 - Added a tutorial for multi-node multi-GPU training with pure PyTorch ([#8071](https://github.com/pyg-team/pytorch_geometric/pull/8071))
 - Added a multinode-multigpu example on `ogbn-papers100M` ([#8070](https://github.com/pyg-team/pytorch_geometric/pull/8070))
diff --git a/examples/kge_fb15k_237.py b/examples/kge_fb15k_237.py
index 036656d27..144074c58 100644
--- a/examples/kge_fb15k_237.py
+++ b/examples/kge_fb15k_237.py
@@ -80,9 +80,10 @@ for epoch in range(1, 501):
     loss = train()
     print(f'Epoch: {epoch:03d}, Loss: {loss:.4f}')
     if epoch % 25 == 0:
-        rank, hits = test(val_data)
+        rank, mrr, hits = test(val_data)
         print(f'Epoch: {epoch:03d}, Val Mean Rank: {rank:.2f}, '
-              f'Val Hits@10: {hits:.4f}')
+              f'Val MRR: {mrr:.4f}, Val Hits@10: {hits:.4f}')
 
-rank, hits_at_10 = test(test_data)
-print(f'Test Mean Rank: {rank:.2f}, Test Hits@10: {hits_at_10:.4f}')
+rank, mrr, hits_at_10 = test(test_data)
+print(f'Test Mean Rank: {rank:.2f}, Test MRR: {mrr:.4f}, '
+      f'Test Hits@10: {hits_at_10:.4f}')
diff --git a/torch_geometric/nn/kge/base.py b/torch_geometric/nn/kge/base.py
index 229fc79aa..63b1c82f0 100644
--- a/torch_geometric/nn/kge/base.py
+++ b/torch_geometric/nn/kge/base.py
@@ -98,9 +98,9 @@ class KGEModel(torch.nn.Module):
         batch_size: int,
         k: int = 10,
         log: bool = True,
-    ) -> Tuple[float, float]:
-        r"""Evaluates the model quality by computing Mean Rank and
-        Hits @ :math:`k` across all possible tail entities.
+    ) -> Tuple[float, float, float]:
+        r"""Evaluates the model quality by computing Mean Rank, MRR and
+        Hits@:math:`k` across all possible tail entities.
 
         Args:
             head_index (torch.Tensor): The head indices.
@@ -115,7 +115,7 @@ class KGEModel(torch.nn.Module):
         arange = range(head_index.numel())
         arange = tqdm(arange) if log else arange
 
-        mean_ranks, hits_at_k = [], []
+        mean_ranks, reciprocal_ranks, hits_at_k = [], [], []
         for i in arange:
             h, r, t = head_index[i], rel_type[i], tail_index[i]
 
@@ -126,12 +126,14 @@ class KGEModel(torch.nn.Module):
             rank = int((torch.cat(scores).argsort(
                 descending=True) == t).nonzero().view(-1))
             mean_ranks.append(rank)
+            reciprocal_ranks.append(1 / (rank + 1))
             hits_at_k.append(rank < k)
 
         mean_rank = float(torch.tensor(mean_ranks, dtype=torch.float).mean())
+        mrr = float(torch.tensor(reciprocal_ranks, dtype=torch.float).mean())
         hits_at_k = int(torch.tensor(hits_at_k).sum()) / len(hits_at_k)
 
-        return mean_rank, hits_at_k
+        return mean_rank, mrr, hits_at_k
 
     @torch.no_grad()
     def random_sample(


Test patch:
diff --git a/test/nn/kge/test_complex.py b/test/nn/kge/test_complex.py
index 0ae36bd11..586587434 100644
--- a/test/nn/kge/test_complex.py
+++ b/test/nn/kge/test_complex.py
@@ -53,6 +53,7 @@ def test_complex():
         loss = model.loss(h, r, t)
         assert loss >= 0.
 
-        mean_rank, hits_at_10 = model.test(h, r, t, batch_size=5, log=False)
-        assert mean_rank <= 10
-        assert hits_at_10 == 1.0
+        mean_rank, mrr, hits = model.test(h, r, t, batch_size=5, log=False)
+        assert 0 <= mean_rank <= 10
+        assert 0 < mrr <= 1
+        assert hits == 1.0
diff --git a/test/nn/kge/test_distmult.py b/test/nn/kge/test_distmult.py
index b9f02f139..df4a20829 100644
--- a/test/nn/kge/test_distmult.py
+++ b/test/nn/kge/test_distmult.py
@@ -19,6 +19,7 @@ def test_distmult():
         loss = model.loss(h, r, t)
         assert loss >= 0.
 
-        mean_rank, hits_at_10 = model.test(h, r, t, batch_size=5, log=False)
-        assert mean_rank <= 10
-        assert hits_at_10 == 1.0
+        mean_rank, mrr, hits = model.test(h, r, t, batch_size=5, log=False)
+        assert 0 <= mean_rank <= 10
+        assert 0 < mrr <= 1
+        assert hits == 1.0
diff --git a/test/nn/kge/test_rotate.py b/test/nn/kge/test_rotate.py
index 87640f386..fd025e0a9 100644
--- a/test/nn/kge/test_rotate.py
+++ b/test/nn/kge/test_rotate.py
@@ -19,6 +19,7 @@ def test_rotate():
         loss = model.loss(h, r, t)
         assert loss >= 0.
 
-        mean_rank, hits_at_10 = model.test(h, r, t, batch_size=5, log=False)
-        assert mean_rank <= 10
-        assert hits_at_10 == 1.0
+        mean_rank, mrr, hits = model.test(h, r, t, batch_size=5, log=False)
+        assert 0 <= mean_rank <= 10
+        assert 0 < mrr <= 1
+        assert hits == 1.0
diff --git a/test/nn/kge/test_transe.py b/test/nn/kge/test_transe.py
index 69c97faa8..13f79c339 100644
--- a/test/nn/kge/test_transe.py
+++ b/test/nn/kge/test_transe.py
@@ -19,6 +19,7 @@ def test_transe():
         loss = model.loss(h, r, t)
         assert loss >= 0.
 
-        mean_rank, hits_at_10 = model.test(h, r, t, batch_size=5, log=False)
-        assert mean_rank <= 10
-        assert hits_at_10 == 1.0
+        mean_rank, mrr, hits = model.test(h, r, t, batch_size=5, log=False)
+        assert 0 <= mean_rank <= 10
+        assert 0 < mrr <= 1
+        assert hits == 1.0