Repository: pyg-team/pytorch_geometric
Pull request #8499 (2023-12-01 09:43:38)

Add link-prediction metric base class and support for `Precision@k`

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8499?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
Attention: `10 lines` in your changes are missing coverage. Please review.
> Comparison is base [(`57fe539`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/57fe53992861be3519c7f3fcde4c7ed0848167a2?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) 88.39% compared to head [(`5b2cbfe`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8499?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) 88.38%.

> :exclamation: Current head 5b2cbfe differs from pull request most recent head f9d9858. Consider uploading reports for the commit f9d9858 to get more accurate results

| [Files](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8499?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) | Patch % | Lines |
|---|---|---|
| [torch\_geometric/nn/metrics.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8499?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL25uL21ldHJpY3MucHk=) | 82.14% | [10 Missing :warning: ](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8499?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) |

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


```diff
@@            Coverage Diff             @@
##           master    #8499      +/-   ##
==========================================
- Coverage   88.39%   88.38%   -0.02%     
==========================================
  Files         478      479       +1     
  Lines       29472    29528      +56     
==========================================
+ Hits        26052    26097      +45     
- Misses       3420     3431      +11     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8499?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 fa6fae298..14f5220a7 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,7 +7,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
-- Added METIS partitioning with CSC/CSR format selection in `ClusterData` ([#8438](https://github.com/pyg-team/pytorch_geometric/pull/8438/))
+- Added support for link-prediction retrieval metrics ([#8499](https://github.com/pyg-team/pytorch_geometric/pull/8499))
+- Added METIS partitioning with CSC/CSR format selection in `ClusterData` ([#8438](https://github.com/pyg-team/pytorch_geometric/pull/8438))
 - Added `is_torch_instance` to check against the original class of compiled models ([#8461](https://github.com/pyg-team/pytorch_geometric/pull/8461))
 - Added dense computation for `AddRandomWalkPE` ([#8431](https://github.com/pyg-team/pytorch_geometric/pull/8431))
 - Added a tutorial for point cloud processing ([#8015](https://github.com/pyg-team/pytorch_geometric/pull/8015))
diff --git a/torch_geometric/nn/metrics.py b/torch_geometric/nn/metrics.py
new file mode 100644
index 000000000..a52bc7985
--- /dev/null
+++ b/torch_geometric/nn/metrics.py
@@ -0,0 +1,142 @@
+from abc import ABC, abstractmethod
+from typing import Optional, Tuple, Union
+
+import torch
+from torch import Tensor
+
+from torch_geometric.utils import scatter
+
+try:
+    import torchmetrics  # noqa
+    WITH_TORCHMETRICS = True
+    BaseMetric = torchmetrics.Metric
+except Exception:
+    WITH_TORCHMETRICS = False
+    BaseMetric = torch.nn.Module
+
+
+class LinkPredMetric(BaseMetric, ABC):
+    r"""An abstract class for computing link prediction retrieval metrics.
+
+    Args:
+        top_k (int): The number of top-:math:`k` predictions to evaluate
+            against.
+    """
+    is_differentiable: Optional[bool] = None
+    higher_is_better: Optional[bool] = None
+    full_state_update: Optional[bool] = None
+
+    def __init__(self, top_k: int):
+        super().__init__()
+
+        if top_k <= 0:
+            raise ValueError(f"'top_k' needs to be a positive integer in "
+                             f"'{self.__class__.__name__}' (got {top_k})")
+
+        self.top_k = top_k
+
+        if WITH_TORCHMETRICS:
+            self.add_state('accum', torch.tensor(0.), dist_reduce_fx='sum')
+            self.add_state('total', torch.tensor(0), dist_reduce_fx='sum')
+        else:
+            self.register_buffer('accum', torch.tensor(0.))
+            self.register_buffer('total', torch.tensor(0))
+
+    def update(
+        self,
+        pred_index_mat: Tensor,
+        edge_label_index: Union[Tensor, Tuple[Tensor, Tensor]],
+    ):
+        r"""Updates the state variables based on the current mini-batch
+        prediction.
+
+        Args:
+            pred_index_mat (torch.Tensor): The top-:math:`k` predictions of
+                every example in the mini-batch with shape
+                :obj:`[batch_size, top_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]`.
+        """
+        if pred_index_mat.size(1) != self.top_k:
+            raise ValueError(f"Expected 'pred_index_mat' to hold {self.top_k} "
+                             f"many indices for every entry "
+                             f"(got {pred_index_mat.size(1)})")
+
+        # Compute a boolean matrix indicating if the k-th prediction is part of
+        # the ground-truth. We do this by flattening both prediction and
+        # target indices, and then determining overlaps via `torch.isin`.
+        max_index = max(
+            pred_index_mat.max() if pred_index_mat.numel() > 0 else 0,
+            edge_label_index[1].max()
+            if edge_label_index[1].numel() > 0 else 0,
+        ) + 1
+        arange = torch.arange(
+            start=0,
+            end=max_index * pred_index_mat.size(0),
+            step=max_index,
+            device=pred_index_mat.device,
+        ).view(-1, 1)
+        flat_pred_index = (pred_index_mat + arange).view(-1)
+        flat_y_index = max_index * edge_label_index[0] + edge_label_index[1]
+
+        pred_isin_mat = torch.isin(flat_pred_index, flat_y_index)
+        pred_isin_mat = pred_isin_mat.view(pred_index_mat.size())
+
+        # Compute the number of targets per example:
+        y_count = scatter(
+            torch.ones_like(edge_label_index[0]),
+            edge_label_index[0],
+            dim=0,
+            dim_size=pred_index_mat.size(0),
+            reduce='sum',
+        )
+
+        metric = self._compute(pred_isin_mat, y_count)
+
+        self.accum += metric.sum()
+        self.total += (y_count > 0).sum()
+
+    def compute(self) -> Tensor:
+        r"""Computes the final metric value."""
+        if self.total == 0:
+            return torch.zeros_like(self.accum)
+        return self.accum / self.total
+
+    def reset(self) -> 'LinkPredMetric':
+        r"""Reset metric state variables to their default value."""
+        if WITH_TORCHMETRICS:
+            super().reset()
+        else:
+            self.accum.zero_()
+            self.total.zero_()
+
+        return self
+
+    @abstractmethod
+    def _compute(self, pred_isin_mat: Tensor, y_count: Tensor) -> Tensor:
+        r"""Compute the specific metric.
+        To be implemented separately for each metric class.
+
+        Args:
+            pred_isin_mat (torch.Tensor): A boolean matrix whose :obj:`(i,k)`
+                element indicates if the :obj:`k`-th prediction for the
+                :obj:`i`-th example is correct or not.
+            y_count (torch.Tensor): A vector indicating the number of
+                ground-truth labels for each example.
+        """
+        raise NotImplementedError
+
+    def __repr__(self) -> str:
+        return f'{self.__class__.__name__}({self.top_k})'
+
+
+class LinkPredPrecision(LinkPredMetric):
+    r"""A link prediction metric to compute Precision@:math`k`.
+
+    Args:
+        top_k (int): The number of top-:math:`k` predictions to evaluate
+            against.
+    """
+    def _compute(self, pred_isin_mat: Tensor, y_count: Tensor) -> Tensor:
+        return pred_isin_mat.sum(dim=-1) / self.top_k


Test patch:
diff --git a/test/nn/test_metrics.py b/test/nn/test_metrics.py
new file mode 100644
index 000000000..a40d5f43e
--- /dev/null
+++ b/test/nn/test_metrics.py
@@ -0,0 +1,48 @@
+from typing import List
+
+import pytest
+import torch
+
+from torch_geometric.nn.metrics import LinkPredPrecision
+
+
+@pytest.mark.parametrize('num_src_nodes', [100])
+@pytest.mark.parametrize('num_dst_nodes', [1000])
+@pytest.mark.parametrize('num_edges', [3000])
+@pytest.mark.parametrize('batch_size', [32])
+@pytest.mark.parametrize('top_k', [1, 10, 100])
+def test_metrics(num_src_nodes, num_dst_nodes, num_edges, batch_size, top_k):
+    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.
+    top_k_pred_mat = pred.topk(top_k, dim=1)[1]
+
+    metric = LinkPredPrecision(top_k)
+    assert str(metric) == f'LinkPredPrecision({top_k})'
+
+    for node_id in torch.split(torch.randperm(num_src_nodes), batch_size):
+        mask = torch.isin(edge_label_index[0], node_id)
+
+        y_batch, y_index = edge_label_index[:, mask]
+        # Remap `y_batch` back to `[0, batch_size - 1]` range:
+        arange = torch.empty(num_src_nodes, dtype=node_id.dtype)
+        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))
+
+    out = metric.compute()
+    metric.reset()
+
+    values: List[float] = []
+    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)
+            precision = float(mask.sum() / top_k)
+            values.append(precision)
+    expected = torch.tensor(values).mean()
+    assert torch.allclose(out, expected)