Add mrr calculation

CHANGELOG.md

@@ -7,7 +7,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 ## [20.07.1] - YYYY-MM-DD
 
 ### Added
-
+- MRR metrics calculation ([#886](https://github.com/catalyst-team/catalyst/pull/886))
 - `CMCScoreCallback` ([#880](https://github.com/catalyst-team/catalyst/pull/880))
 - kornia augmentations `BatchTransformCallback` ([#862](https://github.com/catalyst-team/catalyst/issues/862))
 - `average_precision` and `mean_average_precision` metrics ([#883](https://github.com/catalyst-team/catalyst/pull/883))

catalyst/dl/callbacks/metrics/__init__.py

@@ -26,3 +26,5 @@
     AveragePrecisionCallback,
     MeanAveragePrecisionCallback,
 )
+
+from catalyst.dl.callbacks.metrics.mrr import MRRCallback

catalyst/dl/callbacks/metrics/mrr.py

@@ -0,0 +1,40 @@
+from typing import List
+
+from catalyst.core import MetricCallback
+from catalyst.utils import metrics
+
+
+class MRRCallback(MetricCallback):
+    """Calculates the AUC  per class for each loader.
+
+    .. note::
+        Currently, supports binary and multi-label cases.
+    """
+
+    def __init__(
+        self,
+        input_key: str = "targets",
+        output_key: str = "logits",
+        prefix: str = "mrr",
+        activation: str = "none",
+    ):
+        """
+        Args:
+            input_key (str): input key to use for auc calculation
+                specifies our ``y_true``
+            output_key (str): output key to use for auc calculation;
+                specifies our ``y_pred``
+            prefix (str): name to display for mrr when printing
+            activation (str): An torch.nn activation applied to the outputs.
+                Must be one of ``'none'``, ``'Sigmoid'``, or ``'Softmax2d'``
+        """
+        super().__init__(
+            prefix=prefix,
+            metric_fn=metrics.mrr,
+            input_key=input_key,
+            output_key=output_key,
+            activation=activation,
+        )
+
+
+__all__ = ["MRRCallback"]

catalyst/utils/metrics/__init__.py

@@ -9,3 +9,4 @@
 from .focal import reduced_focal_loss, sigmoid_focal_loss
 from .iou import iou, jaccard
 from .precision import average_precision, mean_average_precision
+from .mrr import mrr

catalyst/utils/metrics/mrr.py

@@ -0,0 +1,42 @@
+"""
+MRR metric.
+"""
+
+import torch
+
+
+def mrr(outputs: torch.Tensor, 
+        targets: torch.Tensor
+    ) -> torch.Tensor:
+
+    """
+    Calculate the MRR score given model ouptputs and targets
+    Args:
+        outputs [batch_size, slate_length] (torch.Tensor): 
+            model outputs, logits
+        targets [batch_szie, slate_length] (torch.Tensor): 
+            ground truth, labels
+    Returns:
+        mrr (float): the mrr score for each slate
+    """
+    max_rank = targets.shape[0]
+
+    _, indices_for_sort = outputs.sort(descending=True, dim=-1)
+    true_sorted_by_preds = torch.gather(targets, dim=-1, index=indices_for_sort)
+    values, indices = torch.max(true_sorted_by_preds, dim=0)
+    indices = indices.type_as(values).unsqueeze(dim=0).t()
+    max_rank_rep = torch.tensor(
+        data=max_rank, device=indices.device, dtype=torch.float32
+    )
+    within_at_mask = (indices < max_rank_rep).type(torch.float32)
+
+    result = torch.tensor(1.0) / (indices + torch.tensor(1.0))
+
+    zero_sum_mask = torch.sum(values) == 0.0
+    result[zero_sum_mask] = 0.0
+
+    mrr = result * within_at_mask
+    return mrr
+
+
+__all__ = ["mrr"]

catalyst/utils/metrics/tests/test_mrr.py

@@ -0,0 +1,39 @@
+import torch
+
+from catalyst.utils import metrics
+
+
+def test_mrr():
+    """
+    Tests for catalyst.utils.metrics.mrr metric.
+    """
+
+    # check 0 simple case
+    y_pred = [0.5, 0.2]
+    y_true = [1.0, 0.0]
+
+    mrr = metrics.mrr(torch.Tensor(y_pred), torch.Tensor(y_true))
+    assert mrr == 1
+
+    # check 1 simple case
+    y_pred = [0.5, 0.2]
+    y_true = [0.0, 1.0]
+
+    mrr = metrics.mrr(torch.Tensor(y_pred), torch.Tensor(y_true))
+    assert mrr == 0.5
+
+    # check 2 simple case
+    y_pred = [0.2, 0.5]
+    y_true = [0.0, 1.0]
+
+    mrr = metrics.mrr(torch.Tensor(y_pred), torch.Tensor(y_true))
+    assert mrr == 1.0
+
+    #test batched slates
+    y_pred_1 = [0.2, 0.5]
+    y_pred_05 = [0.5, 0.2]
+    y_true = [0.0, 1.0]
+
+    mrr = metrics.mrr(torch.Tensor([y_pred_1, y_pred_05]), torch.Tensor([y_true, y_true]))
+    assert mrr[0][0] == 1.0
+    assert mrr[1][0] == 0.5