Logging a class-wise metric with ClasswiseWrapper

[blazejdolicki]

Hey everyone, I am trying to log the metric Dice(average="none") to Neptune using torchmetrics and pytorch lightning. Because this is a class-wise metric, it returns a tensor of shape (num_classes) instead of a scalar which makes it a bit complicated. I found this issue and PR which were supposed to solve the problem, but when I try to use them, I am still getting an error. I would appreciate any hints whether I am doing something wrong in my code or this is just not supported :)

My code looks something like this:

class MyModule(pl.LightningModule):
    def __init__(self):
        (...)
        
        class_labels = ["class_0", "class_1", "class_2"]
        metrics_dict = {"dice_macro": torchmetrics.Dice(average="macro"),
                        "dice_per_class": torchmetrics.ClasswiseWrapper(torchmetrics.Dice(average="none"),
                                                                        labels=class_labels)}
        self.val_metrics = torchmetrics.MetricCollection(metrics_dict)
        
    def validation_step(self, val_batch, batch_idx):
        x, y, w = val_batch
        outputs = self.forward(x)
        loss = self.loss_fn(outputs, y, w)
        self.log("val_loss", loss, on_epoch=True, on_step=False)
        self.val_metrics(outputs, y)
        self.log_dict(self.val_metrics, on_epoch=True, on_step=False)

The error:

Traceback (most recent call last):
    trainer.fit(model, train_loader, val_loader, ckpt_path=checkpoint_dir)
  File "C:\Users\BlazejDolicki\anaconda3\envs\data_science\lib\site-packages\pytorch_lightning\trainer\trainer.py", line 603, in fit
    call._call_and_handle_interrupt(
  File "C:\Users\BlazejDolicki\anaconda3\envs\data_science\lib\site-packages\pytorch_lightning\trainer\call.py", line 38, in _call_and_handle_interrupt
    return trainer_fn(*args, **kwargs)
  File "C:\Users\BlazejDolicki\anaconda3\envs\data_science\lib\site-packages\pytorch_lightning\trainer\trainer.py", line 645, in _fit_impl
    self._run(model, ckpt_path=self.ckpt_path)
  File "C:\Users\BlazejDolicki\anaconda3\envs\data_science\lib\site-packages\pytorch_lightning\trainer\trainer.py", line 1098, in _run
    results = self._run_stage()
  File "C:\Users\BlazejDolicki\anaconda3\envs\data_science\lib\site-packages\pytorch_lightning\trainer\trainer.py", line 1177, in _run_stage
    self._run_train()
  File "C:\Users\BlazejDolicki\anaconda3\envs\data_science\lib\site-packages\pytorch_lightning\trainer\trainer.py", line 1190, in _run_train
    self._run_sanity_check()
  File "C:\Users\BlazejDolicki\anaconda3\envs\data_science\lib\site-packages\pytorch_lightning\trainer\trainer.py", line 1262, in _run_sanity_check
    val_loop.run()
  File "C:\Users\BlazejDolicki\anaconda3\envs\data_science\lib\site-packages\pytorch_lightning\loops\loop.py", line 206, in run
    output = self.on_run_end()
  File "C:\Users\BlazejDolicki\anaconda3\envs\data_science\lib\site-packages\pytorch_lightning\loops\dataloader\evaluation_loop.py", line 184, in on_run_end
    self._on_evaluation_epoch_end()
  File "C:\Users\BlazejDolicki\anaconda3\envs\data_science\lib\site-packages\pytorch_lightning\loops\dataloader\evaluation_loop.py", line 296, in _on_evaluation_epoch_end
    self.trainer._logger_connector.on_epoch_end()
  File "C:\Users\BlazejDolicki\anaconda3\envs\data_science\lib\site-packages\pytorch_lightning\trainer\connectors\logger_connector\logger_connector.py", line 179, in on_epoch_end
    metrics = self.metrics
  File "C:\Users\BlazejDolicki\anaconda3\envs\data_science\lib\site-packages\pytorch_lightning\trainer\connectors\logger_connector\logger_connector.py", line 224, in metrics
    return self.trainer._results.metrics(on_step)
  File "C:\Users\BlazejDolicki\anaconda3\envs\data_science\lib\site-packages\pytorch_lightning\trainer\connectors\logger_connector\result.py", line 584, in metrics
    value = apply_to_collection(result_metric, _ResultMetric, self._get_cache, on_step, include_none=False)
  File "C:\Users\BlazejDolicki\anaconda3\envs\data_science\lib\site-packages\lightning_utilities\core\apply_func.py", line 47, in apply_to_collection
    return function(data, *args, **kwargs)
  File "C:\Users\BlazejDolicki\anaconda3\envs\data_science\lib\site-packages\pytorch_lightning\trainer\connectors\logger_connector\result.py", line 550, in _get_cache
    raise ValueError(
ValueError: The `.compute()` return of the metric logged as 'val_dice_per_class' must be a tensor. Found {'dice_class_0': tensor(0., device='cuda:0'), 'dice_class_1': tensor(0., device='cuda:0'), 'dice_class_2': tensor(0., device='cuda:0')

[patrontheo]

same here

[patrontheo]

@Borda any hint ?