classification metrics (#4043)

* docs + precision + recall + f_beta + refactor

Co-authored-by: Teddy Koker <[email protected]>

* rebase

Co-authored-by: Teddy Koker <[email protected]>

* fixes

Co-authored-by: Teddy Koker <[email protected]>

* added missing file

* docs

* docs

* extra import

Co-authored-by: Teddy Koker <[email protected]>

docs/source/metrics.rst

@@ -126,8 +126,8 @@ Example implementation:
     from pytorch_lightning.metrics import Metric
 
     class MyAccuracy(Metric):
-        def __init__(self, ddp_sync_on_step=False):
-            super().__init__(ddp_sync_on_step=ddp_sync_on_step)
+        def __init__(self, dist_sync_on_step=False):
+            super().__init__(dist_sync_on_step=dist_sync_on_step)
 
             self.add_state("correct", default=torch.tensor(0), dist_reduce_fx="sum")
             self.add_state("total", default=torch.tensor(0), dist_reduce_fx="sum")
@@ -157,6 +157,24 @@ Accuracy
 .. autoclass:: pytorch_lightning.metrics.classification.Accuracy
     :noindex:
 
+Precision
+^^^^^^^^^
+
+.. autoclass:: pytorch_lightning.metrics.classification.Precision
+    :noindex:
+
+Recall
+^^^^^^
+
+.. autoclass:: pytorch_lightning.metrics.classification.Recall
+    :noindex:
+
+Fbeta
+^^^^^
+
+.. autoclass:: pytorch_lightning.metrics.classification.Fbeta
+    :noindex:
+
 Regression Metrics
 ------------------
 

pytorch_lightning/metrics/__init__.py

@@ -1,6 +1,12 @@
 from pytorch_lightning.metrics.metric import Metric
 
-from pytorch_lightning.metrics.classification.accuracy import Accuracy
+from pytorch_lightning.metrics.classification import (
+    Accuracy,
+    Precision,
+    Recall,
+    Fbeta
+)
+
 from pytorch_lightning.metrics.regression import (
     MeanSquaredError,
     MeanAbsoluteError,

pytorch_lightning/metrics/classification/__init__.py

@@ -1 +1,3 @@
 from pytorch_lightning.metrics.classification.accuracy import Accuracy
+from pytorch_lightning.metrics.classification.precision_recall import Precision, Recall
+from pytorch_lightning.metrics.classification.f_beta import Fbeta

pytorch_lightning/metrics/classification/accuracy.py

@@ -31,7 +31,7 @@ class Accuracy(Metric):
             Threshold value for binary or multi-label logits. default: 0.5
         compute_on_step:
             Forward only calls ``update()`` and return None if this is set to False. default: True
-        ddp_sync_on_step:
+        dist_sync_on_step:
             Synchronize metric state across processes at each ``forward()``
             before returning the value at the step. default: False
         process_group:
@@ -52,12 +52,12 @@ def __init__(
         self,
         threshold: float = 0.5,
         compute_on_step: bool = True,
-        ddp_sync_on_step: bool = False,
+        dist_sync_on_step: bool = False,
         process_group: Optional[Any] = None,
     ):
         super().__init__(
             compute_on_step=compute_on_step,
-            ddp_sync_on_step=ddp_sync_on_step,
+            dist_sync_on_step=dist_sync_on_step,
             process_group=process_group,
         )
 
@@ -79,7 +79,6 @@ def _input_format(self, preds: torch.Tensor, target: torch.Tensor):
         if len(preds.shape) == len(target.shape) and preds.dtype == torch.float:
             # binary or multilabel probablities
             preds = (preds >= self.threshold).long()
-
         return preds, target
 
     def update(self, preds: torch.Tensor, target: torch.Tensor):

pytorch_lightning/metrics/classification/f_beta.py

@@ -0,0 +1,119 @@
+import math
+import functools
+from abc import ABC, abstractmethod
+from typing import Any, Callable, Optional, Union
+from collections.abc import Mapping, Sequence
+from collections import namedtuple
+
+import torch
+from torch import nn
+from pytorch_lightning.metrics.metric import Metric
+from pytorch_lightning.metrics.classification.precision_recall import _input_format
+from pytorch_lightning.metrics.utils import METRIC_EPS
+
+
+class Fbeta(Metric):
+    """
+    Computes f_beta metric.
+
+    Works with binary, multiclass, and multilabel data.
+    Accepts logits from a model output or integer class values in prediction.
+    Works with multi-dimensional preds and target.
+
+    Forward accepts
+
+    - ``preds`` (float or long tensor): ``(N, ...)`` or ``(N, C, ...)`` where C is the number of classes
+    - ``target`` (long tensor): ``(N, ...)``
+
+    If preds and target are the same shape and preds is a float tensor, we use the ``self.threshold`` argument.
+    This is the case for binary and multi-label logits.
+
+    If preds has an extra dimension as in the case of multi-class scores we perform an argmax on ``dim=1``.
+
+    Args:
+        num_classes: Number of classes in the dataset.
+        beta: Beta coefficient in the F measure.
+        threshold:
+            Threshold value for binary or multi-label logits. default: 0.5
+
+        average:
+            * `'micro'` computes metric globally
+            * `'macro'` computes metric for each class and then takes the mean
+
+        multilabel: If predictions are from multilabel classification.
+        compute_on_step:
+            Forward only calls ``update()`` and return None if this is set to False. default: True
+        dist_sync_on_step:
+            Synchronize metric state across processes at each ``forward()``
+            before returning the value at the step. default: False
+        process_group:
+            Specify the process group on which synchronization is called. default: None (which selects the entire world)
+
+    Example:
+
+        >>> from pytorch_lightning.metrics import Fbeta
+        >>> target = torch.tensor([0, 1, 2, 0, 1, 2])
+        >>> preds = torch.tensor([0, 2, 1, 0, 0, 1])
+        >>> f_beta = Fbeta(num_classes=3, beta=0.5)
+        >>> f_beta(preds, target)
+        tensor(0.3333)
+
+    """
+    def __init__(
+        self,
+        num_classes: int = 1,
+        beta: float = 1.,
+        threshold: float = 0.5,
+        average: str = 'micro',
+        multilabel: bool = False,
+        compute_on_step: bool = True,
+        dist_sync_on_step: bool = False,
+        process_group: Optional[Any] = None,
+    ):
+        super().__init__(
+            compute_on_step=compute_on_step,
+            dist_sync_on_step=dist_sync_on_step,
+            process_group=process_group,
+        )
+
+        self.num_classes = num_classes
+        self.beta = beta
+        self.threshold = threshold
+        self.average = average
+        self.multilabel = multilabel
+
+        assert self.average in ('micro', 'macro'), \
+            "average passed to the function must be either `micro` or `macro`"
+
+        self.add_state("true_positives", default=torch.zeros(num_classes), dist_reduce_fx="sum")
+        self.add_state("predicted_positives", default=torch.zeros(num_classes), dist_reduce_fx="sum")
+        self.add_state("actual_positives", default=torch.zeros(num_classes), dist_reduce_fx="sum")
+
+    def update(self, preds: torch.Tensor, target: torch.Tensor):
+        """
+        Update state with predictions and targets.
+
+        Args:
+            preds: Predictions from model
+            target: Ground truth values
+        """
+        preds, target = _input_format(self.num_classes, preds, target, self.threshold, self.multilabel)
+
+        self.true_positives += torch.sum(preds * target, dim=1)
+        self.predicted_positives += torch.sum(preds, dim=1)
+        self.actual_positives += torch.sum(target, dim=1)
+
+    def compute(self):
+        """
+        Computes accuracy over state.
+        """
+        if self.average == 'micro':
+            precision = self.true_positives.sum().float() / (self.predicted_positives.sum() + METRIC_EPS)
+            recall = self.true_positives.sum().float() / (self.actual_positives.sum() + METRIC_EPS)
+
+            return (1 + self.beta ** 2) * (precision * recall) / (self.beta ** 2 * precision + recall)
+        elif self.average == 'macro':
+            precision = self.true_positives.float() / (self.predicted_positives + METRIC_EPS)
+            recall = self.true_positives.float() / (self.actual_positives + METRIC_EPS)
+
+            return ((1 + self.beta ** 2) * (precision * recall) / (self.beta ** 2 * precision + recall)).mean()