log_dict: support MetricCollection containing ClasswiseWrapper

[adamjstewart]

What does this PR do?

This PR adds support for logging a MetricCollection containing a ClasswiseWrapper. For example, the following use case is enabled:

from lightning.pytorch import LightningModule
from torchmetrics import MetricCollection
from torchmetrics.wrappers import ClasswiseWrapper
from torchmetrics.classification import MulticlassAccuracy

class MyLightningModule(LightningModule):
    def __init__(self, num_classes):
        ...
        self.val_metrics = MetricCollection({
            'OverallAccuracy': MulticlassAccuracy(average='micro', num_classes=num_classes),
            'AverageAccuracy': MulticlassAccuracy(average='macro', num_classes=num_classes),
            'ClasswiseAccuracy': ClasswiseWrapper(MulticlassAccuracy(average='none', num_classes=num_classes),
        })

    def validation_step(self, batch, batch_idx, dataloader_idx):
        ...
        self.val_metrics(preds, target)
        self.log_dict(self.val_metrics)

I suspect this is a common use case given the multiple issues opened requesting this feature.

Fixes Lightning-AI/torchmetrics#2683 @robmarkcole @SkafteNicki @DimitrisMantas
Fixes Lightning-AI/torchmetrics#2091 @FrenchKrab @SkafteNicki @Northo

It's so common in fact that it's made it to the top of the list of common pitfalls: https://lightning.ai/docs/torchmetrics/stable/pages/lightning.html#common-pitfalls

Luckily, it's also a really easy feature to implement via recursion. Only remaining bug I'm hitting is that lightning.pytorch.trainer.connectors.logger_connector.result._ResultCollection._getcache doesn't like nested dictionaries. It may need similar treatment to support this.

Before submitting

• Was this discussed/agreed via a GitHub issue? (not for typos and docs)
• Did you read the contributor guideline, Pull Request section?
• Did you make sure your PR does only one thing, instead of bundling different changes together?
• Did you make sure to update the documentation with your changes? (if necessary)
• Did you write any new necessary tests? (not for typos and docs)
• Did you verify new and existing tests pass locally with your changes?
• Did you list all the breaking changes introduced by this pull request?
• Did you update the CHANGELOG? (not for typos, docs, test updates, or minor internal changes/refactors)

PR review

Anyone in the community is welcome to review the PR.
Before you start reviewing, make sure you have read the review guidelines. In short, see the following bullet-list:

Reviewer checklist

• Is this pull request ready for review? (if not, please submit in draft mode)
• Check that all items from Before submitting are resolved
• Make sure the title is self-explanatory and the description concisely explains the PR
• Add labels and milestones (and optionally projects) to the PR so it can be classified

---

📚 Documentation preview 📚: https://pytorch-lightning--21507.org.readthedocs.build/en/21507/