Added the ndcg metric [WIP]

ignite/metrics/__init__.py

@@ -19,6 +19,7 @@
 from ignite.metrics.precision import Precision
 from ignite.metrics.psnr import PSNR
 from ignite.metrics.recall import Recall
+from ignite.metrics.recsys.ndcg import NDCG
 from ignite.metrics.root_mean_squared_error import RootMeanSquaredError
 from ignite.metrics.running_average import RunningAverage
 from ignite.metrics.ssim import SSIM

ignite/metrics/recsys/__init__.py

@@ -0,0 +1,5 @@
+from ignite.metrics.recsys.ndcg import NDCG
+
+__all__ = [
+    "NDCG",
+]

ignite/metrics/recsys/ndcg.py

@@ -0,0 +1,73 @@
+from typing import Callable, Optional, Sequence, Union
+
+import torch
+
+from ignite.exceptions import NotComputableError
+from ignite.metrics.metric import Metric
+
+__all__ = ["NDCG"]
+
+
+class NDCG(Metric):
+    def __init__(
+        self,
+        output_transform: Callable = lambda x: x,
+        device: Union[str, torch.device] = torch.device("cpu"),
+        k: Optional[int] = None,
+        log_base: Union[int, float] = 2,
+        exponential: bool = False,
+    ):
+
+        self.log_base = log_base
+        self.k = k
+        self.exponential = exponential
+        super(NDCG, self).__init__(output_transform=output_transform, device=device)
+
+    def _dcg_sample_scores(
+        self, y_true: torch.Tensor, y_score: torch.Tensor, k: Optional[int] = None, log_base: Union[int, float] = 2
+    ) -> torch.Tensor:
+
+        discount = (
+            torch.div(torch.log(torch.arange(y_true.shape[1]) + 2), torch.log(torch.tensor(log_base)))
+            .pow(-1)
+            .to(self._device)
+        )
+        if k is not None:
+            discount[k:] = 0.0
+
+        ranking = torch.argsort(y_score, descending=True)
+        ranked = y_true[torch.arange(ranking.shape[0]).reshape(-1, 1), ranking].to(self._device)
+
+        if self.exponential:
+            ranked = ranked.pow(2) - 1
+
+        discounted_gains = torch.mm(ranked, discount.reshape(-1, 1))
+        return discounted_gains
+
+    def _ndcg_sample_scores(
+        self, y_true: torch.Tensor, y_score: torch.Tensor, k: Optional[int] = None, log_base: Union[int, float] = 2
+    ) -> torch.Tensor:
+
+        gain = self._dcg_sample_scores(y_true, y_score, k, log_base=log_base)
+        normalizing_gain = self._dcg_sample_scores(y_true, y_true, k, log_base=log_base)
+        all_relevant = normalizing_gain != 0
+        normalized_gain = torch.div(gain[all_relevant], normalizing_gain[all_relevant])
+        return normalized_gain
+
+    def reset(self) -> None:
+
+        self.num_examples = 0
+        self.ndcg = torch.tensor(0.0, device=self._device)
+
+    def update(self, output: Sequence[torch.Tensor]) -> None:
+
+        y, y_pred = output[0].detach(), output[1].detach()
+        gain = self._ndcg_sample_scores(y, y_pred, k=self.k, log_base=self.log_base)
+        self.ndcg = torch.add(self.ndcg, torch.sum(gain))
+        self.num_examples += y_pred.shape[0]
+
+    def compute(self) -> float:
+        if self.num_examples == 0:
+            raise NotComputableError("NGCD must have at least one example before it can be computed.")
+
+        return (self.ndcg / self.num_examples).item()

tests/ignite/metrics/test_ndcg.py

@@ -0,0 +1,67 @@
+import numpy as np
+import pytest
+import torch
+from sklearn.metrics import ndcg_score
+
+from ignite.exceptions import NotComputableError
+from ignite.metrics.recsys.ndcg import NDCG
+
+
+@pytest.mark.parametrize(
+    "true_score, pred_score",
+    [
+        (torch.tensor([[1.0, 2.0, 3.0, 4.0, 5.0]]), torch.tensor([[0.1, 0.2, 0.3, 0.4, 0.5]])),
+        (torch.tensor([[2.9, 5.6, 3.8, 7.9, 6.2]]), torch.tensor([[3.7, 4.8, 3.9, 4.3, 4.9]])),
+        (
+            torch.tensor([[1.0, 2.0, 3.0, 4.0, 5.0], [1.2, 4.5, 8.9, 5.6, 7.2], [2.9, 5.6, 3.8, 7.9, 6.2]]),
+            torch.tensor([[0.1, 0.2, 0.3, 0.4, 0.5], [3.7, 4.8, 3.9, 4.3, 4.9], [3.7, 4.8, 3.9, 4.3, 4.9]]),
+        ),
+    ],
+)
+@pytest.mark.parametrize("k", [None, 2, 3])
+def test_output_cpu(true_score, pred_score, k):
+
+    device = "cpu"
+
+    ndcg = NDCG(k=k, device=device)
+    ndcg.update([true_score, pred_score])
+    result_ignite = ndcg.compute()
+    result_sklearn = ndcg_score(true_score.numpy(), pred_score.numpy(), k=k)
+
+    np.testing.assert_allclose(np.array(result_ignite), result_sklearn, rtol=2e-7)
+
+
+@pytest.mark.parametrize(
+    "true_score, pred_score",
+    [
+        (torch.tensor([[1.0, 2.0, 3.0, 4.0, 5.0]]), torch.tensor([[0.1, 0.2, 0.3, 0.4, 0.5]])),
+        (torch.tensor([[2.9, 5.6, 3.8, 7.9, 6.2]]), torch.tensor([[3.7, 4.8, 3.9, 4.3, 4.9]])),
+        (
+            torch.tensor([[1.0, 2.0, 3.0, 4.0, 5.0], [1.2, 4.5, 8.9, 5.6, 7.2], [2.9, 5.6, 3.8, 7.9, 6.2]]),
+            torch.tensor([[0.1, 0.2, 0.3, 0.4, 0.5], [3.7, 4.8, 3.9, 4.3, 4.9], [3.7, 4.8, 3.9, 4.3, 4.9]]),
+        ),
+    ],
+)
+@pytest.mark.parametrize("k", [None, 2, 3])
+@pytest.mark.skipif(torch.cuda.device_count() < 1, reason="Skip if no GPU")
+def test_output_gpu(true_score, pred_score, k):
+
+    device = "cuda"
+    ndcg = NDCG(k=k, device=device)
+    ndcg.update([true_score, pred_score])
+    result_ignite = ndcg.compute()
+    result_sklearn = ndcg_score(true_score.numpy(), pred_score.numpy(), k=k)
+
+    np.testing.assert_allclose(np.array(result_ignite), result_sklearn, rtol=2e-7)
+
+
+def test_reset():
+
+    true_score = torch.tensor([[1.0, 2.0, 3.0, 4.0, 5.0]])
+    pred_score = torch.tensor([[0.1, 0.2, 0.3, 0.4, 0.5]])
+    ndcg = NDCG()
+    ndcg.update([true_score, pred_score])
+    ndcg.reset()
+
+    with pytest.raises(NotComputableError, match=r"NGCD must have at least one example before it can be computed."):
+        ndcg.compute()