[Metrics] Class reduction similar to sklearn

CHANGELOG.md

@@ -11,6 +11,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 - Added hooks to metric module interface ([#2528](https://github.com/PyTorchLightning/pytorch-lightning/pull/2528/))
 
+- Added `class_reduction` similar to sklearn for classification metrics ([#3322](https://github.com/PyTorchLightning/pytorch-lightning/pull/3322))
+
 ### Changed
 
 
@@ -125,7 +127,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Fixed adding val step argument to metrics ([#2986](https://github.com/PyTorchLightning/pytorch-lightning/pull/2986))
 - Fixed an issue that caused `Trainer.test()` to stall in ddp mode ([#2997](https://github.com/PyTorchLightning/pytorch-lightning/pull/2997))
 - Fixed gathering of results with tensors of varying shape ([#3020](https://github.com/PyTorchLightning/pytorch-lightning/pull/3020))
-- Fixed batch size auto-scaling feature to set the new value on the correct model attribute ([#3043](https://github.com/PyTorchLightning/pytorch-lightning/pull/3043)) 
+- Fixed batch size auto-scaling feature to set the new value on the correct model attribute ([#3043](https://github.com/PyTorchLightning/pytorch-lightning/pull/3043))
 - Fixed automatic batch scaling not working with half precision ([#3045](https://github.com/PyTorchLightning/pytorch-lightning/pull/3045))
 - Fixed setting device to root gpu ([#3042](https://github.com/PyTorchLightning/pytorch-lightning/pull/3042))
 

pytorch_lightning/metrics/classification.py

@@ -52,26 +52,28 @@ class Accuracy(TensorMetric):
     def __init__(
             self,
             num_classes: Optional[int] = None,
-            reduction: str = 'elementwise_mean',
+            class_reduction: str = 'micro',
             reduce_group: Any = None,
             reduce_op: Any = None,
     ):
         """
         Args:
             num_classes: number of classes
-            reduction: a method to reduce metric score over labels (default: takes the mean)
-                Available reduction methods:
-                - elementwise_mean: takes the mean
-                - none: pass array
-                - sum: add elements
+            class_reduction: reduction method for multiclass problems
+
+                - ``'micro'``: calculate metrics globally (default)
+                - ``'macro'``: calculate metrics for each label, and find their unweighted mean.
+                - ``'weighted'``: calculate metrics for each label, and find their unweighted mean.
+                - ``'none'``: returns calculated metric per class
+
             reduce_group: the process group to reduce metric results from DDP
             reduce_op: the operation to perform for ddp reduction
         """
         super().__init__(name='accuracy',
                          reduce_group=reduce_group,
                          reduce_op=reduce_op)
         self.num_classes = num_classes
-        self.reduction = reduction
+        self.class_reduction = class_reduction
 
     def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
         """
@@ -85,7 +87,7 @@ def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
             A Tensor with the classification score.
         """
         return accuracy(pred=pred, target=target,
-                        num_classes=self.num_classes, reduction=self.reduction)
+                        num_classes=self.num_classes, class_reduction=self.class_reduction)
 
 
 class ConfusionMatrix(TensorMetric):
@@ -146,13 +148,10 @@ class PrecisionRecallCurve(TensorCollectionMetric):
         >>> pred = torch.tensor([0, 1, 2, 3])
         >>> target = torch.tensor([0, 1, 2, 2])
         >>> metric = PrecisionRecallCurve()
-        >>> prec, recall, thr = metric(pred, target)
-        >>> prec
-        tensor([0.3333, 0.0000, 0.0000, 1.0000])
-        >>> recall
-        tensor([1., 0., 0., 0.])
-        >>> thr
-        tensor([1., 2., 3.])
+        >>> metric(pred, target)   # doctest: +NORMALIZE_WHITESPACE
+        (tensor([0.3333, 0.0000, 0.0000, 1.0000]),
+         tensor([1., 0., 0., 0.]),
+         tensor([1., 2., 3.]))
 
     """
 
@@ -206,7 +205,7 @@ class Precision(TensorMetric):
 
         >>> pred = torch.tensor([0, 1, 2, 3])
         >>> target = torch.tensor([0, 1, 2, 2])
-        >>> metric = Precision(num_classes=4)
+        >>> metric = Precision(num_classes=4, class_reduction='macro')
         >>> metric(pred, target)
         tensor(0.7500)
 
@@ -215,26 +214,28 @@ class Precision(TensorMetric):
     def __init__(
             self,
             num_classes: Optional[int] = None,
-            reduction: str = 'elementwise_mean',
+            class_reduction: str = 'micro',
             reduce_group: Any = None,
             reduce_op: Any = None,
     ):
         """
         Args:
             num_classes: number of classes
-            reduction: a method to reduce metric score over labels (default: takes the mean)
-                Available reduction methods:
-                - elementwise_mean: takes the mean
-                - none: pass array
-                - sum: add elements
+            class_reduction: reduction method for multiclass problems
+
+                - ``'micro'``: calculate metrics globally (default)
+                - ``'macro'``: calculate metrics for each label, and find their unweighted mean.
+                - ``'weighted'``: calculate metrics for each label, and find their unweighted mean.
+                - ``'none'``: returns calculated metric per class
+
             reduce_group: the process group to reduce metric results from DDP
             reduce_op: the operation to perform for ddp reduction
         """
         super().__init__(name='precision',
                          reduce_group=reduce_group,
                          reduce_op=reduce_op)
         self.num_classes = num_classes
-        self.reduction = reduction
+        self.class_reduction = class_reduction
 
     def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
         """
@@ -249,7 +250,7 @@ def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
         """
         return precision(pred=pred, target=target,
                          num_classes=self.num_classes,
-                         reduction=self.reduction)
+                         class_reduction=self.class_reduction)
 
 
 class Recall(TensorMetric):
@@ -262,25 +263,27 @@ class Recall(TensorMetric):
         >>> target = torch.tensor([0, 1, 2, 2])
         >>> metric = Recall()
         >>> metric(pred, target)
-        tensor(0.6250)
+        tensor(0.7500)
 
     """
 
     def __init__(
             self,
             num_classes: Optional[int] = None,
-            reduction: str = 'elementwise_mean',
+            class_reduction: str = 'micro',
             reduce_group: Any = None,
             reduce_op: Any = None,
     ):
         """
         Args:
             num_classes: number of classes
-            reduction: a method to reduce metric score over labels (default: takes the mean)
-                Available reduction methods:
-                - elementwise_mean: takes the mean
-                - none: pass array
-                - sum: add elements
+            class_reduction: reduction method for multiclass problems
+
+                - ``'micro'``: calculate metrics globally (default)
+                - ``'macro'``: calculate metrics for each label, and find their unweighted mean.
+                - ``'weighted'``: calculate metrics for each label, and find their unweighted mean.
+                - ``'none'``: returns calculated metric per class
+
             reduce_group: the process group to reduce metric results from DDP
             reduce_op: the operation to perform for ddp reduction
         """
@@ -289,7 +292,7 @@ def __init__(
                          reduce_op=reduce_op)
 
         self.num_classes = num_classes
-        self.reduction = reduction
+        self.class_reduction = class_reduction
 
     def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
         """
@@ -305,7 +308,7 @@ def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
         return recall(pred=pred,
                       target=target,
                       num_classes=self.num_classes,
-                      reduction=self.reduction)
+                      class_reduction=self.class_reduction)
 
 
 class AveragePrecision(TensorMetric):
@@ -425,7 +428,7 @@ class FBeta(TensorMetric):
 
         >>> pred = torch.tensor([0, 1, 2, 3])
         >>> target = torch.tensor([0, 1, 2, 2])
-        >>> metric = FBeta(0.25)
+        >>> metric = FBeta(0.25, class_reduction='macro')
         >>> metric(pred, target)
         tensor(0.7361)
     """
@@ -434,19 +437,21 @@ def __init__(
             self,
             beta: float,
             num_classes: Optional[int] = None,
-            reduction: str = 'elementwise_mean',
+            class_reduction: str = 'micro',
             reduce_group: Any = None,
             reduce_op: Any = None,
     ):
         """
         Args:
             beta: determines the weight of recall in the combined score.
             num_classes: number of classes
-            reduction: a method to reduce metric score over labels (default: takes the mean)
-                Available reduction methods:
-                - elementwise_mean: takes the mean
-                - none: pass array
-                - sum: add elements
+            class_reduction: reduction method for multiclass problems
+
+                - ``'micro'``: calculate metrics globally (default)
+                - ``'macro'``: calculate metrics for each label, and find their unweighted mean.
+                - ``'weighted'``: calculate metrics for each label, and find their unweighted mean.
+                - ``'none'``: returns calculated metric per class
+
             reduce_group: the process group to reduce metric results from DDP
             reduce_op: the operation to perform for DDP reduction
         """
@@ -456,7 +461,7 @@ def __init__(
 
         self.beta = beta
         self.num_classes = num_classes
-        self.reduction = reduction
+        self.class_reduction = class_reduction
 
     def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
         """
@@ -471,7 +476,7 @@ def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
         """
         return fbeta_score(pred=pred, target=target,
                            beta=self.beta, num_classes=self.num_classes,
-                           reduction=self.reduction)
+                           class_reduction=self.class_reduction)
 
 
 class F1(TensorMetric):
@@ -483,26 +488,28 @@ class F1(TensorMetric):
 
         >>> pred = torch.tensor([0, 1, 2, 3])
         >>> target = torch.tensor([0, 1, 2, 2])
-        >>> metric = F1()
+        >>> metric = F1(class_reduction='macro')
         >>> metric(pred, target)
         tensor(0.6667)
     """
 
     def __init__(
             self,
             num_classes: Optional[int] = None,
-            reduction: str = 'elementwise_mean',
+            class_reduction: str = 'micro',
             reduce_group: Any = None,
             reduce_op: Any = None,
     ):
         """
         Args:
             num_classes: number of classes
-            reduction: a method to reduce metric score over labels (default: takes the mean)
-                Available reduction methods:
-                - elementwise_mean: takes the mean
-                - none: pass array
-                - sum: add elements
+            class_reduction: reduction method for multiclass problems
+
+                - ``'micro'``: calculate metrics globally (default)
+                - ``'macro'``: calculate metrics for each label, and find their unweighted mean.
+                - ``'weighted'``: calculate metrics for each label, and find their unweighted mean.
+                - ``'none'``: returns calculated metric per class
+
             reduce_group: the process group to reduce metric results from DDP
             reduce_op: the operation to perform for ddp reduction
         """
@@ -511,7 +518,7 @@ def __init__(
                          reduce_op=reduce_op)
 
         self.num_classes = num_classes
-        self.reduction = reduction
+        self.class_reduction = class_reduction
 
     def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
         """
@@ -526,7 +533,7 @@ def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
         """
         return f1_score(pred=pred, target=target,
                         num_classes=self.num_classes,
-                        reduction=self.reduction)
+                        class_reduction=self.class_reduction)
 
 
 class ROC(TensorCollectionMetric):
@@ -538,13 +545,10 @@ class ROC(TensorCollectionMetric):
         >>> pred = torch.tensor([0, 1, 2, 3])
         >>> target = torch.tensor([0, 1, 2, 2])
         >>> metric = ROC()
-        >>> fps, tps, thresholds = metric(pred, target)
-        >>> fps
-        tensor([0.0000, 0.3333, 0.6667, 0.6667, 1.0000])
-        >>> tps
-        tensor([0., 0., 0., 1., 1.])
-        >>> thresholds
-        tensor([4., 3., 2., 1., 0.])
+        >>> metric(pred, target)   # doctest: +NORMALIZE_WHITESPACE
+        (tensor([0.0000, 0.3333, 0.6667, 0.6667, 1.0000]),
+         tensor([0., 0., 0., 1., 1.]),
+         tensor([4., 3., 2., 1., 0.]))
 
     """