True update() for MetricCollection

[IgorHoholko]

🚀 Feature

MetricCollection class overrides update method of parent nn.ModuleDict class. Is it possible to add update from parent's class to MetricCollection?

 def update_(self, metrics: Union[List[Metric], Tuple[Metric], Dict[str, Metric]]):
        # ... another work for List[Metric], Tuple[Metric] ...
        else:
            return super().update(metrics)

Example of using

metrics = MetricCollection({metric_name: getattr(metrics_obj, metric_name)
                                    for metric_name in self.metrics})

self.train_metrics = metrics.clone(prefix='train_')
self.val_metrics = metrics.clone(prefix='val_')
self.test_metrics = metrics.clone(prefix='test_')
self.test_metrics.update({metric_name: getattr(metrics_obj, metric_name)
                            for metric_name in  ('auroc', 'auroc_macro', 'auroc_weighted')})

[SkafteNicki]

You can use super to call the parent method. Basic example

import torch
from torchmetrics import MetricCollection, Accuracy, Precision, Recall, ConfusionMatrix
target = torch.tensor([0, 2, 0, 2, 0, 1, 0, 2])
preds = torch.tensor([2, 1, 2, 0, 1, 2, 2, 2])
metrics = MetricCollection([
    Accuracy(),
    Precision(num_classes=3, average='macro'),
    Recall(num_classes=3, average='macro')
])
print(metrics)
>>>MetricCollection(
>>>  (Accuracy): Accuracy()
>>>  (Precision): Precision()
>>>  (Recall): Recall()
>>>)
# call super
super(type(metrics), metrics).update({'confmat': ConfusionMatrix(num_classes=3)})
print(metrics)
>>>MetricCollection(
>>>  (Accuracy): Accuracy()
>>>  (Precision): Precision()
>>>  (Recall): Recall()
>>>  (confmat): ConfusionMatrix()
>>>)

[IgorHoholko]

@SkafteNicki Makes sence. Thanks!

[SkafteNicki]

@IgorHoholko you are welcome. Closing the issue. Feel free to re-open if necessary :]

[maximsch2]

Should we do something like

def add_metrics(self, metrics):
   super().update(metrics)

?

[IgorHoholko]

@maximsch2 It would definitely be helpful.

[IgorHoholko]

metrics.add_metrics() looks better and more intuitive than super(type(metrics), metrics).update() :)

[SkafteNicki]

fine by me :]
@IgorHoholko want to send a PR?

[IgorHoholko]

@SkafteNicki Yes, a bit later.