Logging metrics that return dicts

[Aceticia]

I have a torchmetric class that calculates many different aspects of the predictions. More specifically, it returns a dict of Dict[MetricName, float]. It looks like self.log and self.log_dict only take single output metrics? Is there a way of directly logging the torchmetric object?

For now I'm just updating it, taking its return dict and call self.log_dict. Works, but not as clean

[tshu-w]

In fact, the best practice is to pass the metric object to self.log/log_dict directly, so we don't need to reset the metric object manually.

class MyModel(LightningModule):

    def __init__(self):
        ...
        self.accuracy = torchmetrics.Accuracy()

    def training_step(self, batch, batch_idx):
        x, y = batch
        preds = self(x)
        ...
        # log step metric
        self.accuracy(preds, y)
        self.log('train_acc_step', self.accuracy)
        ...

    def training_epoch_end(self, outs):
        # log epoch metric
        self.log('train_acc_epoch', self.accuracy)

For more details, please refers to https://torchmetrics.readthedocs.io/en/stable/pages/lightning.html

[Aceticia]

That is what I'm trying to do, but it looks like the self.log would only take one metric object that returns one value? My metric returns a dict, which contains multiple values. It wouldn't make sense to log them all as the same metric.

[tshu-w]

Then I suggest using MetricCollection instead of a single metric to return dict?