Repository: Lightning-AI/torchmetrics

Issue:
Precision Loss in ConfusionMatrix
## 🐛 Bug

In Pytorch, `torch.zeros` and `.float()` methods create or convert a tensor to be a `torch.float32` tensor. Compare to the default `np.float64` data type in NumPy, it cannot provide enough precision. In `ConfusionMatrix`, the default values are torch.float32, which leads to wrong calculation (although maybe very little) against ground truth.

I think we only need to change `torch.zeros(num_classes, num_classes)` to `torch.zeros(num_classes, num_classes, dtype=torch.float64)`. I believe this change introduces negligible memory cost.

### To Reproduce

```python
>>> torch.tensor(90170814.).item()
90170816.0
```

### Expected Behaviour

```python
>>> torch.tensor(90170814., dtype=torch.float64).item()
90170814.0
```

### Additional context

`ConfusionMatrix` is the base of many other metrics, such as `IoU`. I found this problem by comparing the results to another NumPy-based metric in a semantic segmentation task using `IoU` as the metric. I think in today's deep learning dataset, float32 precision is not enough for the implementation of metrics.

Based on the above observation, maybe all `PyTorchLightning/metrics` should be implemented based on `torch.float64`




Discussion:
Hi! thanks for your contribution!, great first issue!
Considering we may not need an accurate metric when validating (I see that even 16-precision is supported), maybe a flag to enable 64-precision in testing is a better choice.

Resolving patch:
diff --git a/torchmetrics/functional/classification/cohen_kappa.py b/torchmetrics/functional/classification/cohen_kappa.py
--- a/torchmetrics/functional/classification/cohen_kappa.py
+++ b/torchmetrics/functional/classification/cohen_kappa.py
@@ -23,6 +23,7 @@
 
 def _cohen_kappa_compute(confmat: Tensor, weights: Optional[str] = None) -> Tensor:
     confmat = _confusion_matrix_compute(confmat)
+    confmat = confmat.float() if not confmat.is_floating_point() else confmat
     n_classes = confmat.shape[0]
     sum0 = confmat.sum(dim=0, keepdim=True)
     sum1 = confmat.sum(dim=1, keepdim=True)
diff --git a/torchmetrics/functional/classification/confusion_matrix.py b/torchmetrics/functional/classification/confusion_matrix.py
--- a/torchmetrics/functional/classification/confusion_matrix.py
+++ b/torchmetrics/functional/classification/confusion_matrix.py
@@ -28,7 +28,6 @@ def _confusion_matrix_update(
     if mode not in (DataType.BINARY, DataType.MULTILABEL):
         preds = preds.argmax(dim=1)
         target = target.argmax(dim=1)
-
     if multilabel:
         unique_mapping = ((2 * target + preds) + 4 * torch.arange(num_classes, device=preds.device)).flatten()
         minlength = 4 * num_classes
@@ -48,8 +47,8 @@ def _confusion_matrix_compute(confmat: Tensor, normalize: Optional[str] = None)
     allowed_normalize = ('true', 'pred', 'all', 'none', None)
     assert normalize in allowed_normalize, \
         f"Argument average needs to one of the following: {allowed_normalize}"
-    confmat = confmat.float()
     if normalize is not None and normalize != 'none':
+        confmat = confmat.float() if not confmat.is_floating_point() else confmat
         cm = None
         if normalize == 'true':
             cm = confmat / confmat.sum(axis=1, keepdim=True)
diff --git a/torchmetrics/metric.py b/torchmetrics/metric.py
--- a/torchmetrics/metric.py
+++ b/torchmetrics/metric.py
@@ -312,8 +312,13 @@ def _apply(self, fn):
         to the correct device when `.to`, `.cuda`, etc methods are called
         """
         this = super()._apply(fn)
-        # Also apply fn to metric states
-        for key in this._defaults.keys():
+        # Also apply fn to metric states and defaults
+        for key, value in this._defaults.items():
+            if isinstance(value, Tensor):
+                this._defaults[key] = fn(value)
+            elif isinstance(value, Sequence):
+                this._defaults[key] = [fn(v) for v in value]
+
             current_val = getattr(this, key)
             if isinstance(current_val, Tensor):
                 setattr(this, key, fn(current_val))


Test patch:
diff --git a/tests/bases/test_metric.py b/tests/bases/test_metric.py
--- a/tests/bases/test_metric.py
+++ b/tests/bases/test_metric.py
@@ -269,9 +269,13 @@ def test_device_and_dtype_transfer(tmpdir):
 
     metric = metric.double()
     assert metric.x.dtype == torch.float64
+    metric.reset()
+    assert metric.x.dtype == torch.float64
 
     metric = metric.half()
     assert metric.x.dtype == torch.float16
+    metric.reset()
+    assert metric.x.dtype == torch.float16
 
 
 def test_warning_on_compute_before_update():