New modular metric interface

pytorch_lightning/metrics/converters.py

@@ -63,28 +63,32 @@ def new_func(*args, **kwargs):
     return decorator_fn
 
 
-def _convert_to_tensor(data: Any) -> Any:
+def convert_to_tensor(data: Any, dtype=None, device=None) -> Any:
     """
     Maps all kind of collections and numbers to tensors.
 
     Args:
         data: the data to convert to tensor
 
+        dtype: data type to convert to
+
+        device: device to cast to
+
     Return:
         the converted data
     """
     if isinstance(data, numbers.Number):
-        return torch.tensor([data])
+        return torch.tensor([data], dtype=dtype, device=device)
     # is not array of object
     elif isinstance(data, np.ndarray) and np_str_obj_array_pattern.search(data.dtype.str) is None:
-        return torch.from_numpy(data)
+        return torch.from_numpy(data).to(device=device, dtype=dtype)
     elif isinstance(data, torch.Tensor):
-        return data
+        return data.to(device=device, dtype=dtype)
 
     raise TypeError(f"The given type ('{type(data).__name__}') cannot be converted to a tensor!")
 
 
-def _convert_to_numpy(data: Union[torch.Tensor, np.ndarray, numbers.Number]) -> np.ndarray:
+def convert_to_numpy(data: Union[torch.Tensor, np.ndarray, numbers.Number]) -> np.ndarray:
     """Convert all tensors and numpy arrays to numpy arrays.
 
     Args:
@@ -114,7 +118,7 @@ def _numpy_metric_input_conversion(func_to_decorate: Callable) -> Callable:
         Callable: the decorated function
     """
     return _apply_to_inputs(
-        apply_to_collection, (torch.Tensor, np.ndarray, numbers.Number), _convert_to_numpy)(func_to_decorate)
+        apply_to_collection, (torch.Tensor, np.ndarray, numbers.Number), convert_to_numpy)(func_to_decorate)
 
 
 def _tensor_metric_output_conversion(func_to_decorate: Callable) -> Callable:
@@ -127,7 +131,7 @@ def _tensor_metric_output_conversion(func_to_decorate: Callable) -> Callable:
     Return:
         Callable: the decorated function
     """
-    return _apply_to_outputs(_convert_to_tensor)(func_to_decorate)
+    return _apply_to_outputs(convert_to_tensor)(func_to_decorate)
 
 
 def _numpy_metric_conversion(func_to_decorate: Callable) -> Callable:
@@ -161,7 +165,7 @@ def _tensor_metric_input_conversion(func_to_decorate: Callable) -> Callable:
         Callable: the decorated function
     """
     return _apply_to_inputs(
-        apply_to_collection, (torch.Tensor, np.ndarray, numbers.Number), _convert_to_tensor)(func_to_decorate)
+        apply_to_collection, (torch.Tensor, np.ndarray, numbers.Number), convert_to_tensor)(func_to_decorate)
 
 
 def _tensor_collection_metric_output_conversion(func_to_decorate: Callable) -> Callable:
@@ -175,7 +179,7 @@ def _tensor_collection_metric_output_conversion(func_to_decorate: Callable) -> C
         Callable: the decorated function
     """
     return _apply_to_outputs(apply_to_collection, (torch.Tensor, np.ndarray, numbers.Number),
-                             _convert_to_tensor)(func_to_decorate)
+                             convert_to_tensor)(func_to_decorate)
 
 
 def _tensor_metric_conversion(func_to_decorate: Callable) -> Callable:
@@ -215,10 +219,11 @@ def _tensor_collection_metric_conversion(func_to_decorate: Callable) -> Callable
     return _tensor_collection_metric_output_conversion(func_convert_inputs)
 
 
-def _sync_ddp_if_available(result: Union[torch.Tensor],
-                           group: Optional[Any] = None,
-                           reduce_op: Optional[torch.distributed.ReduceOp] = None,
-                           ) -> torch.Tensor:
+def sync_ddp_if_available(result: Union[torch.Tensor],
+                          group: Optional[Any] = None,
+                          reduce_op: Optional[Any] = None,
+                          ddp_normalize=False,
+                          ) -> torch.Tensor:
     """
     Function to reduce the tensors from several ddp processes to one master process
 
@@ -243,11 +248,14 @@ def _sync_ddp_if_available(result: Union[torch.Tensor],
         torch.distributed.all_reduce(result, op=reduce_op, group=group,
                                      async_op=False)
 
+        if ddp_normalize:
+            result / torch.distributed.get_world_size(group)
+
     return result
 
 
 def sync_ddp(group: Optional[Any] = None,
-             reduce_op: Optional[torch.distributed.ReduceOp] = None) -> Callable:
+             reduce_op: Optional[Any] = None) -> Callable:
     """
     This decorator syncs a functions outputs across different processes for DDP.
 
@@ -262,14 +270,14 @@ def sync_ddp(group: Optional[Any] = None,
 
     def decorator_fn(func_to_decorate):
         return _apply_to_outputs(apply_to_collection, torch.Tensor,
-                                 _sync_ddp_if_available, group=group,
+                                 sync_ddp_if_available, group=group,
                                  reduce_op=reduce_op)(func_to_decorate)
 
     return decorator_fn
 
 
 def numpy_metric(group: Optional[Any] = None,
-                 reduce_op: Optional[torch.distributed.ReduceOp] = None) -> Callable:
+                 reduce_op: Optional[Any] = None) -> Callable:
     """
     This decorator shall be used on all function metrics working on numpy arrays.
     It handles the argument conversion and DDP reduction for metrics working on numpy.
@@ -292,7 +300,7 @@ def decorator_fn(func_to_decorate):
 
 
 def tensor_metric(group: Optional[Any] = None,
-                  reduce_op: Optional[torch.distributed.ReduceOp] = None) -> Callable:
+                  reduce_op: Optional[Any] = None) -> Callable:
     """
     This decorator shall be used on all function metrics working on tensors.
     It handles the argument conversion and DDP reduction for metrics working on tensors.
@@ -314,7 +322,7 @@ def decorator_fn(func_to_decorate):
 
 
 def tensor_collection_metric(group: Optional[Any] = None,
-                             reduce_op: Optional[torch.distributed.ReduceOp] = None) -> Callable:
+                             reduce_op: Optional[Any] = None) -> Callable:
     """
     This decorator shall be used on all function metrics working on tensors and returning collections
     that cannot be converted to tensors.

pytorch_lightning/metrics/functional/classification.py

@@ -256,7 +256,7 @@ def confusion_matrix(
     """
     num_classes = get_num_classes(pred, target, None)
 
-    unique_labels = target.view(-1) * num_classes + pred.view(-1)
+    unique_labels = (target.view(-1) * num_classes + pred.view(-1)).to(torch.int)
 
     bins = torch.bincount(unique_labels, minlength=num_classes ** 2)
     cm = bins.reshape(num_classes, num_classes).squeeze().float()

pytorch_lightning/metrics/metric.py

@@ -1,11 +1,13 @@
 from abc import ABC, abstractmethod
 from typing import Any, Optional
+import numbers
 
 import torch
 import torch.distributed
+import numpy as np
 
 from pytorch_lightning.metrics.converters import (
-    tensor_metric, numpy_metric, tensor_collection_metric)
+    sync_ddp_if_available, convert_to_tensor, convert_to_numpy)
 from pytorch_lightning.utilities.apply_func import apply_to_collection
 from pytorch_lightning.utilities.device_dtype_mixin import DeviceDtypeModuleMixin
 
@@ -17,6 +19,16 @@ class Metric(DeviceDtypeModuleMixin, torch.nn.Module, ABC):
     Should be used to implement metrics that
     1. Return multiple Outputs
     2. Handle their own DDP sync
+
+    Metric hooks that can be implemented are:
+        input_convert: pre-forward hook that takes care of input conversion
+        output_convert: post-forward hook that takes care of output convertion
+        ddp_sync: implementation of ddp sync
+        compute: post-ddp sync for additional metric computations
+
+    Call order:
+        input_convert -> forward -> output_convert -> ddp_sync -> compute
+
     """
 
     def __init__(self, name: str):
@@ -29,18 +41,51 @@ def __init__(self, name: str):
         self.name = name
         self._dtype = torch.get_default_dtype()
         self._device = torch.device('cpu')
+        self.register_forward_pre_hook(self.input_convert)
+        self.register_forward_hook(self.output_convert)
+        self.register_forward_hook(self.ddp_sync)
+        self.register_forward_hook(self.compute)
 
     @abstractmethod
     def forward(self, *args, **kwargs) -> torch.Tensor:
         """
         Implements the actual metric computation.
 
         Returns:
-            metric value
+            metric value or metric state
 
         """
         raise NotImplementedError
 
+    def compute(self, module, input, output) -> torch.Tensor:
+        """
+        Implement additionally metric computations to be done after the ddp sync
+
+        Args:
+            module: current metric module
+
+            input: input to forward method
+
+            output: output from forward method
+
+        Returns:
+            final metric value
+
+        """
+        return output
+
+    def ddp_sync(self, module, input, output):
+        """
+
+        """
+        return output
+
+    def input_convert(self, module, input):
+        return input
+
+    def output_convert(self, module, input, output):
+        return output
+
 
 class TensorMetric(Metric):
     """
@@ -51,7 +96,8 @@ class TensorMetric(Metric):
 
     def __init__(self, name: str,
                  reduce_group: Optional[Any] = None,
-                 reduce_op: Optional[Any] = None):
+                 reduce_op: Optional[Any] = None,
+                 ddp_normalize: bool = False):
         """
 
         Args:
@@ -62,15 +108,23 @@ def __init__(self, name: str,
                 Defaults to sum.
         """
         super().__init__(name)
-        self._orig_call = tensor_metric(group=reduce_group,
-                                        reduce_op=reduce_op)(super().__call__)
+        self.reduce_group = reduce_group
+        self.reduce_op = reduce_op
+        self.ddp_normalize = ddp_normalize
+
+    def input_convert(self, module, input):
+        return apply_to_collection(input,
+                                   (torch.Tensor, np.ndarray, numbers.Number),
+                                   convert_to_tensor,
+                                   self.dtype, self.device)
 
-    def __call__(self, *args, **kwargs) -> torch.Tensor:
-        def _to_device_dtype(x: torch.Tensor) -> torch.Tensor:
-            return x.to(device=self.device, dtype=self.dtype, non_blocking=True)
+    def output_convert(self, module, input, output):
+        return apply_to_collection(output, torch.Tensor, convert_to_tensor,
+                                   self.dtype, self.device)
 
-        return apply_to_collection(self._orig_call(*args, **kwargs), torch.Tensor,
-                                   _to_device_dtype)
+    def ddp_sync(self, module, input, output):
+        return apply_to_collection(output, torch.Tensor, sync_ddp_if_available,
+                                   self.reduce_group, self.reduce_op, self.ddp_normalize)
 
 
 class TensorCollectionMetric(Metric):
@@ -92,7 +146,8 @@ class TensorCollectionMetric(Metric):
 
     def __init__(self, name: str,
                  reduce_group: Optional[Any] = None,
-                 reduce_op: Optional[Any] = None):
+                 reduce_op: Optional[Any] = None,
+                 ddp_normalize: bool = False):
         """
 
         Args:
@@ -103,15 +158,25 @@ def __init__(self, name: str,
                 Defaults to sum.
         """
         super().__init__(name)
-        self._orig_call = tensor_collection_metric(group=reduce_group,
-                                                   reduce_op=reduce_op)(super().__call__)
+        self.reduce_group = reduce_group
+        self.reduce_op = reduce_op
+        self.ddp_normalize = ddp_normalize
+
+    def input_convert(self, module, input):
+        return apply_to_collection(input,
+                                   (torch.Tensor, np.ndarray, numbers.Number),
+                                   convert_to_tensor,
+                                   self.dtype, self.device)
 
-    def __call__(self, *args, **kwargs) -> torch.Tensor:
-        def _to_device_dtype(x: torch.Tensor) -> torch.Tensor:
-            return x.to(device=self.device, dtype=self.dtype, non_blocking=True)
+    def output_convert(self, module, input, output):
+        return apply_to_collection(output,
+                                   (torch.Tensor, np.ndarray, numbers.Number),
+                                   convert_to_tensor,
+                                   self.dtype, self.device)
 
-        return apply_to_collection(self._orig_call(*args, **kwargs), torch.Tensor,
-                                   _to_device_dtype)
+    def ddp_sync(self, module, input, output):
+        return apply_to_collection(output, torch.Tensor, sync_ddp_if_available,
+                                   self.reduce_group, self.reduce_op, self.ddp_normalize)
 
 
 class NumpyMetric(Metric):
@@ -124,7 +189,8 @@ class NumpyMetric(Metric):
 
     def __init__(self, name: str,
                  reduce_group: Optional[Any] = None,
-                 reduce_op: Optional[Any] = None):
+                 reduce_op: Optional[Any] = None,
+                 ddp_normalize: bool = False):
         """
 
         Args:
@@ -135,12 +201,21 @@ def __init__(self, name: str,
                 Defaults to sum.
         """
         super().__init__(name)
-        self._orig_call = numpy_metric(group=reduce_group,
-                                       reduce_op=reduce_op)(super().__call__)
-
-    def __call__(self, *args, **kwargs) -> torch.Tensor:
-        def _to_device_dtype(x: torch.Tensor) -> torch.Tensor:
-            return x.to(device=self.device, dtype=self.dtype, non_blocking=True)
-
-        return apply_to_collection(self._orig_call(*args, **kwargs), torch.Tensor,
-                                   _to_device_dtype)
+        self.reduce_group = reduce_group
+        self.reduce_op = reduce_op
+        self.ddp_normalize = ddp_normalize
+
+    def input_convert(self, module, input):
+        return apply_to_collection(input,
+                                   (torch.Tensor, np.ndarray, numbers.Number),
+                                   convert_to_numpy)
+
+    def output_convert(self, module, input, output):
+        return apply_to_collection(output,
+                                   (torch.Tensor, np.ndarray, numbers.Number),
+                                   convert_to_tensor,
+                                   self.dtype, self.device)
+
+    def ddp_sync(self, module, input, output):
+        return apply_to_collection(output, torch.Tensor, sync_ddp_if_available,
+                                   self.reduce_group, self.reduce_op, self.ddp_normalize)

tests/metrics/test_classification.py

@@ -45,7 +45,6 @@ def test_confusion_matrix(normalize):
 
     target = (torch.arange(120) % 3).view(-1, 1)
     pred = target.clone()
-
     cm = conf_matrix(pred, target)
     assert isinstance(cm, torch.Tensor)