Repository: Lightning-AI/torchmetrics

Issue:
MultiTaskWrapper doesn't support being logged as a metric
## 🚀 Feature

Make MultiTaskWrapper (and possibly other wrappers) support logging in pytorch_lightning.LightningModule. I expected it would just work with self.log_dict, but it doesn't, since it's subclassed from Metric. It would be great if it had MetricCollection-like functionality.

### Motivation

Better compatibility with pytorch_lightning is always good; also, this standardizes the APIs of different torchmetrics tools.

### Pitch

TorchMetrics's class metrics are great at standardizing metric calculation, which in most scenarios allows to share boilerplate code, such as logging metric in each step (see Additional context for an example), between modules. This stems from torchmetrics' compatibility with pytorch_lightning. Probably, most people using this library, use it in conjunction with lightning, so increasing compatibility with lightning will definitely be an improvement.

### Alternatives

Of course, one could use MultiTaskWrapper in the way outlined in the official example, calculating the metric, getting a simple python dictionary, logging the dictionary. But this would have a few shortcomings:
- The epoch-wise metrics could be wrong! E.g. if one uses a classification metric with average="macro", logging each batch separately and then averaging over the batches to get the epoch-level metric is just wrong. One would need to manually save all targets and predictions, and calculate the metric in the on_*_end hook.
- This makes room for greater confusion, I personally expected log_dict to work, and only after looking at the source code did I realize that it can't! There is not a word about pytorch_lightning in the docs, so one expects it to work as do the other metrics!

### Additional context

Let me give an example of removing shared boilerplate logic from submodules. This could be the abstract parent class with the boilerplate code, which all submodules share.
```
    def training_step(
        self, batch, batch_idx: int
    ):
        preds, target = self.shared_step(batch, batch_idx)
        train_loss = self.loss(preds, target)
        self.train_metrics(preds, target)

        self.log("train_loss", train_loss, on_epoch=True)
        self.log_dict(self.train_metrics, on_epoch=True) # type: ignore

        return train_loss

    def validation_step(
        self, batch, batch_idx: int
    ):
        preds, target = self.shared_step(batch)
        val_loss = self.loss(preds, target)
        self.val_metrics(preds, target)

        self.log("val_loss", val_loss)
        self.log_dict(self.val_metrics) # type: ignore

    def test_step(
        self, batch, batch_idx: int
    ):
        preds, target = self.shared_step(batch, batch_idx)

        self.test_metrics(preds, target)
        self.log_dict(self.test_metrics, on_step=False, on_epoch=True) # type: ignore
```

And with that all a child class has to do is specify metrics and loss in __init__:
```
        def loss(probs, target):
            return nn.functional.nll_loss(
                torch.log(probs), target, ignore_index=pad_value
            )

        metrics = MetricCollection(
            {
                "AUROC": AUROC(
                    task="multiclass",
                    num_classes=num_types,
                    ignore_index=pad_value,
                    average="macro"
                ),
                "PR-AUC": AveragePrecision(
                    task="multiclass",
                    num_classes=num_types,
                    ignore_index=pad_value,
                    average="macro"
                ),
                "Accuracy": Accuracy(
                    task="multiclass",
                    num_classes=num_types,
                    ignore_index=pad_value,
                    average="macro"
                ),
                "F1Score": F1Score(
                    task="multiclass",
                    num_classes=num_types,
                    ignore_index=pad_value,
                    average="macro"
                ),
            }
        )

        super().__init__(
            loss=loss,
            metrics=metrics,
        )
```
and overload the shared_step to get the specific prediction logic. This would even work with different multitask modules, if only MultiTaskWrapper supported this behaviour.



Discussion:
Hi! thanks for your contribution!, great first issue!

Resolving patch:
diff --git a/src/torchmetrics/wrappers/multitask.py b/src/torchmetrics/wrappers/multitask.py
--- a/src/torchmetrics/wrappers/multitask.py
+++ b/src/torchmetrics/wrappers/multitask.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # this is just a bypass for this module name collision with built-in one
-from typing import Any, Dict, Optional, Sequence, Union
+from typing import Any, Dict, Iterable, Optional, Sequence, Tuple, Union
 
 from torch import Tensor, nn
 
@@ -102,6 +102,18 @@ def __init__(
         super().__init__()
         self.task_metrics = nn.ModuleDict(task_metrics)
 
+    def items(self) -> Iterable[Tuple[str, nn.Module]]:
+        """Iterate over task and task metrics."""
+        return self.task_metrics.items()
+
+    def keys(self) -> Iterable[str]:
+        """Iterate over task names."""
+        return self.task_metrics.keys()
+
+    def values(self) -> Iterable[nn.Module]:
+        """Iterate over task metrics."""
+        return self.task_metrics.values()
+
     @staticmethod
     def _check_task_metrics_type(task_metrics: Dict[str, Union[Metric, MetricCollection]]) -> None:
         if not isinstance(task_metrics, dict):


Test patch:
diff --git a/tests/integrations/test_lightning.py b/tests/integrations/test_lightning.py
--- a/tests/integrations/test_lightning.py
+++ b/tests/integrations/test_lightning.py
@@ -28,6 +28,8 @@
 from torchmetrics import MetricCollection
 from torchmetrics.aggregation import SumMetric
 from torchmetrics.classification import BinaryAccuracy, BinaryAveragePrecision
+from torchmetrics.regression import MeanSquaredError
+from torchmetrics.wrappers import MultitaskWrapper
 
 from integrations.helpers import no_warning_call
 from integrations.lightning.boring_model import BoringModel
@@ -358,6 +360,52 @@ def on_train_epoch_end(self):
     assert torch.allclose(tensor(logged["DiffMetric_epoch"]), model.diff, atol=2e-4)
 
 
+def test_task_wrapper_lightning_logging(tmpdir):
+    """Test that MultiTaskWrapper works with Lightning modules."""
+
+    class TestModel(BoringModel):
+        def __init__(self) -> None:
+            super().__init__()
+            self.metric = MultitaskWrapper({"classification": BinaryAccuracy(), "regression": MeanSquaredError()})
+            self.accuracy = BinaryAccuracy()
+            self.mse = MeanSquaredError()
+
+        def training_step(self, batch, batch_idx):
+            preds = torch.rand(10)
+            target = torch.rand(10)
+            self.metric(
+                {"classification": preds.round(), "regression": preds},
+                {"classification": target.round(), "regression": target},
+            )
+            self.accuracy(preds.round(), target.round())
+            self.mse(preds, target)
+            self.log("accuracy", self.accuracy, on_epoch=True)
+            self.log("mse", self.mse, on_epoch=True)
+            self.log_dict(self.metric, on_epoch=True)
+            return self.step(batch)
+
+    model = TestModel()
+
+    trainer = Trainer(
+        default_root_dir=tmpdir,
+        limit_train_batches=2,
+        limit_val_batches=0,
+        max_epochs=1,
+        log_every_n_steps=1,
+    )
+    with no_warning_call(
+        UserWarning,
+        match="Torchmetrics v0.9 introduced a new argument class property called.*",
+    ):
+        trainer.fit(model)
+
+    logged = trainer.logged_metrics
+    assert torch.allclose(logged["accuracy_step"], logged["classification_step"])
+    assert torch.allclose(logged["accuracy_epoch"], logged["classification_epoch"])
+    assert torch.allclose(logged["mse_step"], logged["regression_step"])
+    assert torch.allclose(logged["mse_epoch"], logged["regression_epoch"])
+
+
 def test_scriptable(tmpdir):
     """Test that lightning modules can still be scripted even if metrics cannot."""