Logging a class-wise metric with ClasswiseWrapper #1462
-
|
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: The error: |
Beta Was this translation helpful? Give feedback.
Replies: 3 comments 2 replies
-
|
same here |
Beta Was this translation helpful? Give feedback.
-
|
@Borda any hint ? |
Beta Was this translation helpful? Give feedback.
-
|
For those stumbling upon this now, there is a comment on this here: https://lightning.ai/docs/torchmetrics/stable/pages/lightning.html#common-pitfalls Essentially log it with the manual approach, rather than the metric object. class MyModule(LightningModule):
def __init__(self):
super().__init__()
self.train_metrics = MetricCollection(
{
"macro_accuracy": MinMaxMetric(MulticlassAccuracy(num_classes=5, average="macro")),
"weighted_accuracy": MinMaxMetric(MulticlassAccuracy(num_classes=5, average="weighted")),
},
prefix="train_",
)
def training_step(self, batch, batch_idx):
...
# logging the MetricCollection object directly will fail
self.log_dict(self.train_metrics(preds, target))
# manually computing the result and then logging will work
batch_values = self.train_metrics(preds, target)
self.log_dict(batch_values, on_step=True, on_epoch=False)
...
def on_train_epoch_end(self):
self.train_metrics.reset() |
Beta Was this translation helpful? Give feedback.
-
|
@Northo's workaround above is on the right track — manually computing and logging the dict works. Let me add a bit more context on why and provide a clean pattern you can reuse. Why Recommended pattern (TorchMetrics >= v1.9.0): Use from torchmetrics import MetricCollection
from torchmetrics.wrappers import ClasswiseWrapper
from torchmetrics.classification import MulticlassF1Score, MulticlassAccuracy
class MyModel(L.LightningModule):
def __init__(self, num_classes=5):
super().__init__()
self.metrics = MetricCollection({
"f1": ClasswiseWrapper(
MulticlassF1Score(num_classes=num_classes, average="none"),
labels=["cat", "dog", "bird", "fish", "horse"],
prefix="f1_",
),
"acc": MulticlassAccuracy(num_classes=num_classes, average="macro"),
})
def validation_step(self, batch, batch_idx):
x, y = batch
logits = self(x)
# MetricCollection.forward() returns a flat dict
metric_dict = self.metrics(logits, y)
# log_dict handles the flat dict natively
self.log_dict(metric_dict, on_epoch=True, prog_bar=False)This gives you keys like If you need manual control (e.g. custom logging names): def on_validation_epoch_end(self):
result = self.metrics.compute()
for name, value in result.items():
self.log(name, value)
self.metrics.reset()Docs: ClasswiseWrapper | MetricCollection |
Beta Was this translation helpful? Give feedback.
-
|
Thanks for the insight, @Borda! |
Beta Was this translation helpful? Give feedback.
@Northo's workaround above is on the right track — manually computing and logging the dict works. Let me add a bit more context on why and provide a clean pattern you can reuse.
Why
self.logdoesn't work directly withClasswiseWrapper:Lightning's
self.logexpects either a scalarTensoror aMetricwhose.compute()returns a scalar.ClasswiseWrapper.compute()returns adict[str, Tensor], so Lightning doesn't know how to handle it. This isn't a bug per se — it's a design mismatch between Lightning's logging and per-class metrics.Recommended pattern (TorchMetrics >= v1.9.0):
Use
ClasswiseWrapperinside aMetricCollection, then log the whole collection withself.log_dict: