Repository: Lightning-AI/torchmetrics (language: python)

Issue:
MinMaxMetric not working in Pytorch Lightning
## 🐛 Bug

When using the `MinMaxMetric` wrapper in a LightningModule, as shown in Torchmetrics [docs](https://lightning.ai/docs/torchmetrics/stable/pages/lightning.html), it generates a ValueError: `ValueError: The '.compute()' return of the metric logged as 'val/macro_jaccard' must be a tensor. Found {'raw': tensor(0.3260, device='cuda:0'), 'max': tensor(0.3260, device='cuda:0'), 'min': tensor(0.3231, device='cuda:0')}`

### To Reproduce
Here is a sample code to reproduce the error
<details>
  <summary>Code sample</summary>

```python
import torch
import torch.nn.functional as F
import pytorch_lightning as pl
from torchmetrics import JaccardIndex, MetricCollection
from torchmetrics.wrappers import MinMaxMetric

class LitModel(pl.LightningModule):
    def __init__(self):
        super().__init__()
        # Simple model (1 conv layer for demonstration)
        self.model = torch.nn.Conv2d(3, 2, kernel_size=3, padding=1)
        
        # Define train metrics and wrap them with MinMaxMetric
        self.train_metrics = MetricCollection({
            'weighted_jaccard': MinMaxMetric(JaccardIndex(task='multiclass', num_classes=2, average='weighted')),
            'macro_jaccard': MinMaxMetric(JaccardIndex(task='multiclass', num_classes=2, average='macro'))
        })

        # Clone metrics for validation
        self.val_metrics = self.train_metrics.clone(prefix='val/')

    def forward(self, x):
        return self.model(x)
    
    def step(self, batch):
        x, y = batch
        logits = self(x)
        loss = F.cross_entropy(logits, y)
        return logits, y, loss

    def training_step(self, batch, batch_idx):
        logits, y, loss = self.step(batch)

        # Log loss
        self.log("train/loss", loss, on_step=True, on_epoch=True)

        # Update and log metrics
        self.train_metrics(logits.argmax(dim=1), y)
        self.log_dict(self.train_metrics, on_step=True, on_epoch=True)

        return loss

    def validation_step(self, batch, batch_idx):
        logits, y, loss = self.step(batch)

        # Log validation loss
        self.log("val/loss", loss, on_step=False, on_epoch=True)

        # Update and log validation metrics
        self.val_metrics(logits.argmax(dim=1), y)
        self.log_dict(self.val_metrics, on_step=False, on_epoch=True)

        return loss

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=1e-3)

# Simulated data for training and validation
def generate_fake_data(image_size=(3, 64, 64), num_classes=2):
    x = torch.randn(*image_size)  # Fake single image input
    y = torch.randint(0, num_classes, (image_size[1], image_size[2]))  # Fake segmentation labels
    return x, y

# Create DataLoaders
train_data = [generate_fake_data() for _ in range(100)]  # 100 batches of training data
val_data = [generate_fake_data() for _ in range(20)]     # 20 batches of validation data

train_loader = torch.utils.data.DataLoader(train_data, batch_size=16)
val_loader = torch.utils.data.DataLoader(val_data, batch_size=16)

# Initialize the model
model = LitModel()

# Initialize a PyTorch Lightning Trainer and fit the model
trainer = pl.Trainer(max_epochs=3, log_every_n_steps=10, devices=[0])
trainer.fit(model, train_dataloaders=train_loader, val_dataloaders=val_loader)
```

</details>

### Expected behavior

It should not generate a ValueError and should log 3 metrics: the actual one, the min one, and the max one.

### Environment

- TorchMetrics version (if build from source, add commit SHA): 1.4.2
- Python & PyTorch Version (e.g., 1.0): 3.11.9 & 2.4.0
- Any other relevant information such as OS (e.g., Linux): Ubuntu VM

### Additional context

If the MinMaxMetric is used outside of a MetricCollection, it generates another error. The error doesn't happen at `.compute` time but before, when doing `.log`. Here is the traceback (but maybe I should open a second issue?):
```
Traceback (most recent call last):
  File "<string>", line 1, in <module>
  File "/home/patron/mambaforge/envs/segmentator/lib/python3.11/site-packages/pytorch_lightning/core/module.py", line 608, in log_dict
    for k, v in dictionary.items(**kwargs):
                ^^^^^^^^^^^^^^^^
  File "/home/patron/mambaforge/envs/segmentator/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1729, in __getattr__
    raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
AttributeError: 'MinMaxMetric' object has no attribute 'items'
```



Discussion / hints:

Fix patch:
diff --git a/docs/source/pages/lightning.rst b/docs/source/pages/lightning.rst
index 09ceafb331f..a7678f3e2a5 100644
--- a/docs/source/pages/lightning.rst
+++ b/docs/source/pages/lightning.rst
@@ -108,7 +108,7 @@ also manually log the output of the metrics.
 
     class MyModule(LightningModule):
 
-        def __init__(self):
+        def __init__(self, num_classes):
             ...
             self.train_acc = torchmetrics.classification.Accuracy(task="multiclass", num_classes=num_classes)
             self.valid_acc = torchmetrics.classification.Accuracy(task="multiclass", num_classes=num_classes)
@@ -157,6 +157,43 @@ Additionally, we highly recommend that the two ways of logging are not mixed as
                 self.valid_acc.update(logits, y)
                 self.log('valid_acc', self.valid_acc, on_step=True, on_epoch=True)
 
+In general if you are logging multiple metrics we highly recommend that you combine them into a single metric object
+using the :class:`~torchmetrics.MetricCollection` class and then replacing the ``self.log`` calls with ``self.log_dict``,
+assuming that all metrics receive the same input.
+
+.. testcode:: python
+
+    class MyModule(LightningModule):
+
+        def __init__(self):
+            ...
+            self.train_metrics = torchmetrics.MetricCollection(
+                {
+                    "accuracy": torchmetrics.classification.Accuracy(task="multiclass", num_classes=num_classes),
+                    "f1": torchmetrics.classification.F1(task="multiclass", num_classes=num_classes),
+                },
+                prefix="train_",
+            )
+            self.valid_metrics = self.train_metrics.clone(prefix="valid_")
+
+        def training_step(self, batch, batch_idx):
+            x, y = batch
+            preds = self(x)
+            ...
+            batch_value = self.train_metrics(preds, y)
+            self.log_dict(batch_value)
+
+        def on_train_epoch_end(self):
+            self.train_metrics.reset()
+
+        def validation_step(self, batch, batch_idx):
+            logits = self(x)
+            ...
+            self.valid_metrics.update(logits, y)
+
+        def on_validation_epoch_end(self, outputs):
+            self.log_dict(self.valid_metrics.compute())
+            self.valid_metrics.reset()
 
 ***************
 Common Pitfalls
@@ -164,6 +201,41 @@ Common Pitfalls
 
 The following contains a list of pitfalls to be aware of:
 
+* Logging a `MetricCollection` object directly using ``self.log_dict`` is only supported if all metrics in the
+  collection returns a scalar tensor. If any of the metrics in the collection returns a non-scalar tensor,
+  the logging will fail. This can especially happen when either nesting multiple ``MetricCollection`` objects or when
+  using wrapper metrics such as :class:`~torchmetrics.wrappers.ClasswiseWrapper`,
+  :class:`~torchmetrics.wrappers.MinMaxMetric` etc. inside a ``MetricCollection`` since all these wrappers return
+  dicts or lists of tensors. It is still possible to log such nested metrics manually because the ``MetricCollection``
+  object will try to flatten everything into a single dict. Example:
+
+.. testcode:: python
+
+    class MyModule(LightningModule):
+
+        def __init__(self):
+            super().__init__()
+            self.train_metrics = MetricCollection(
+                {
+                    "macro_accuracy": MinMaxMetric(MulticlassAccuracy(num_classes=5, average="macro")),
+                    "weighted_accuracy": MinMaxMetric(MulticlassAccuracy(num_classes=5, average="weighted")),
+                },
+                prefix="train_",
+            )
+
+        def training_step(self, batch, batch_idx):
+            ...
+            # logging the MetricCollection object directly will fail
+            self.log_dict(self.train_metrics(preds, target))
+
+            # manually computing the result and then logging will work
+            batch_values = self.train_metrics(preds, target)
+            self.log_dict(batch_values, on_step=True, on_epoch=False)
+            ...
+
+        def on_train_epoch_end(self):
+            self.train_metrics.reset()
+
 * Modular metrics contain internal states that should belong to only one DataLoader. In case you are using multiple DataLoaders,
   it is recommended to initialize a separate modular metric instances for each DataLoader and use them separately. The same holds
   for using separate metrics for training, validation and testing.


Test patch:
diff --git a/tests/integrations/test_lightning.py b/tests/integrations/test_lightning.py
index 513c351e2a0..302a4353f76 100644
--- a/tests/integrations/test_lightning.py
+++ b/tests/integrations/test_lightning.py
@@ -30,7 +30,7 @@
 from torchmetrics.classification import BinaryAccuracy, BinaryAveragePrecision, MulticlassAccuracy
 from torchmetrics.regression import MeanAbsoluteError, MeanSquaredError
 from torchmetrics.utilities.prints import rank_zero_only
-from torchmetrics.wrappers import ClasswiseWrapper, MultitaskWrapper
+from torchmetrics.wrappers import ClasswiseWrapper, MinMaxMetric, MultitaskWrapper
 
 from integrations.lightning.boring_model import BoringModel
 
@@ -523,35 +523,26 @@ def __init__(self) -> None:
                 },
                 prefix="train_",
             )
-
-            self.val_metrics = MetricCollection(
-                {
-                    "macro_accuracy": MulticlassAccuracy(num_classes=5, average="macro"),
-                    "classwise_accuracy": ClasswiseWrapper(MulticlassAccuracy(num_classes=5, average=None)),
-                },
-                prefix="val_",
-            )
+            self.val_metrics = self.train_metrics.clone(prefix="val_")
 
         def training_step(self, batch, batch_idx):
             loss = self(batch).sum()
-
             preds = torch.randint(0, 5, (100,), device=batch.device)
             target = torch.randint(0, 5, (100,), device=batch.device)
 
             self.train_metrics.update(preds, target)
             batch_values = self.train_metrics.compute()
             self.log_dict(batch_values, on_step=True, on_epoch=False)
-
             return {"loss": loss}
 
         def validation_step(self, batch, batch_idx):
             preds = torch.randint(0, 5, (100,), device=batch.device)
             target = torch.randint(0, 5, (100,), device=batch.device)
-
             self.val_metrics.update(preds, target)
 
         def on_validation_epoch_end(self):
             self.log_dict(self.val_metrics.compute(), on_step=False, on_epoch=True)
+            self.val_metrics.reset()
 
     model = TestModel()
 
@@ -559,7 +550,7 @@ def on_validation_epoch_end(self):
         default_root_dir=tmpdir,
         limit_train_batches=2,
         limit_val_batches=2,
-        max_epochs=1,
+        max_epochs=2,
         log_every_n_steps=1,
     )
     trainer.fit(model)
@@ -572,3 +563,61 @@ def on_validation_epoch_end(self):
     for i in range(5):
         assert f"train_multiclassaccuracy_{i}" in logged
         assert f"val_multiclassaccuracy_{i}" in logged
+
+
+def test_collection_minmax_lightning_integration(tmpdir):
+    """Check the integration of MinMaxWrapper, MetricCollection and LightningModule.
+
+    See issue: https://github.com/Lightning-AI/torchmetrics/issues/2763
+
+    """
+
+    class TestModel(BoringModel):
+        def __init__(self) -> None:
+            super().__init__()
+            self.train_metrics = MetricCollection(
+                {
+                    "macro_accuracy": MinMaxMetric(MulticlassAccuracy(num_classes=5, average="macro")),
+                    "weighted_accuracy": MinMaxMetric(MulticlassAccuracy(num_classes=5, average="weighted")),
+                },
+                prefix="train_",
+            )
+            self.val_metrics = self.train_metrics.clone(prefix="val_")
+
+        def training_step(self, batch, batch_idx):
+            loss = self(batch).sum()
+            preds = torch.randint(0, 5, (100,), device=batch.device)
+            target = torch.randint(0, 5, (100,), device=batch.device)
+
+            self.train_metrics.update(preds, target)
+            batch_values = self.train_metrics.compute()
+            self.log_dict(batch_values, on_step=True, on_epoch=False)
+            return {"loss": loss}
+
+        def validation_step(self, batch, batch_idx):
+            preds = torch.randint(0, 5, (100,), device=batch.device)
+            target = torch.randint(0, 5, (100,), device=batch.device)
+            self.val_metrics.update(preds, target)
+
+        def on_validation_epoch_end(self):
+            self.log_dict(self.val_metrics.compute(), on_step=False, on_epoch=True)
+            self.val_metrics.reset()
+
+    model = TestModel()
+
+    trainer = Trainer(
+        default_root_dir=tmpdir,
+        limit_train_batches=2,
+        limit_val_batches=2,
+        max_epochs=2,
+        log_every_n_steps=1,
+    )
+    trainer.fit(model)
+
+    logged = trainer.logged_metrics
+
+    # check that all metrics are logged
+    for prefix in ["train_", "val_"]:
+        for metric in ["macro_accuracy", "weighted_accuracy"]:
+            for key in ["max", "min", "raw"]:
+                assert f"{prefix}{metric}_{key}" in logged