TorchMetrics Documentation

TorchMetrics is a collection of over 100 PyTorch metric implementations designed for distributed and scalable PyTorch applications.

Key Features

• Standardized Interface: Increases reproducibility across projects.
• Reduced Boilerplate: Simplifies the implementation of common metrics.
• Automatic Accumulation: Handles the accumulation of metric states over multiple batches automatically.
• Distributed Training Optimization: Metrics are optimized for distributed environments.
• Automatic Synchronization: Handles synchronization of metric states between multiple devices (e.g., multiple GPUs).

Integration with PyTorch Lightning

When used with PyTorch Lightning, you get additional benefits:

• Automatic Device Placement: Module metrics are automatically moved to the correct device (CPU/GPU).
• Native Logging: Seamless support for logging metrics within Lightning to further reduce boilerplate.

Some metrics require access to individual batch states (e.g., for ranking or correlation). In these cases, you can initialize states with an empty list as the default value in add_state.

Key requirements for list states:

• Set dist_reduce_fx="cat" in add_state to concatenate tensors across batches/processes.
• In update(), append the batch tensors to the list (e.g., self.preds.append(preds)).
• In compute(), use the dim_zero_cat helper function from torchmetrics.utilities to standardize the list into a single concatenated tensor. This ensures compatibility between distributed and non-distributed modes.

Warning: Calling reset() clears the list state entirely. If you need the data after a reset, you must manually copy it using copy.deepcopy before calling reset().

from torchmetrics import Metric
from torchmetrics.utilities import dim_zero_cat

class MySpearmanCorrCoef(Metric):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.add_state("preds", default=[], dist_reduce_fx="cat")
        self.add_state("target", default=[], dist_reduce_fx="cat")

    def update(self, preds: Tensor, target: Tensor) -> None:
        self.preds.append(preds)
        self.target.append(target)

    def compute(self):
        # Use dim_zero_cat to handle list states correctly in all modes
        preds = dim_zero_cat(self.preds)
        target = dim_zero_cat(self.target)
        # ... computation logic ...
        return corrcoef

Module-based metrics (subclasses of torchmetrics.Metric) are designed to automate state management. They handle:

• Automatic accumulation of metric states over multiple batches.
• Automatic synchronization of states across multiple devices (CPUs, single GPUs, or multi-GPU/multi-node setups).
• Metric arithmetic.

To use a module metric, initialize it, move it to the desired device using .to(device), and call it on your predictions and targets. Use .compute() to retrieve the final accumulated result and .reset() to clear the internal state for a new epoch or evaluation cycle.

import torch
import torchmetrics

# initialize metric
metric = torchmetrics.classification.Accuracy(task="multiclass", num_classes=5)

# move the metric to device
device = "cuda" if torch.cuda.is_available() else "cpu"
metric.to(device)

n_batches = 10
for i in range(n_batches):
    preds = torch.randn(10, 5).softmax(dim=-1).to(device)
    target = torch.randint(5, (10,)).to(device)

    # metric on current batch
    acc = metric(preds, target)
    print(f"Accuracy on batch {i}: {acc}")

# metric on all batches using custom accumulation
acc = metric.compute()
print(f"Accuracy on all data: {acc}")

The recommended way to log metrics in Lightning is to pass the metric object itself to self.log.

When you log a Metric object directly:

1. Lightning uses the on_step and on_epoch flags from self.log(...) to determine logging frequency.
2. If on_epoch=True, Lightning automatically calls .compute() at the end of the epoch to log the final value.

Important Constraints:

• self.log only supports scalar tensors. Metrics that return non-scalar outputs (like ConfusionMatrix, ROC, MeanAveragePrecision, or ROUGEScore) must be handled manually by computing their values and logging the resulting scalars/dictionaries.
• The sync_dist, sync_dist_group, and reduce_fx flags in self.log(...) do not affect Metric objects, as these objects handle their own distributed synchronization logic internally. This does not apply to the functional metric API.

class MyModule(LightningModule):
    def __init__(self, num_classes):
        super().__init__()
        self.train_acc = torchmetrics.classification.Accuracy(task="multiclass", num_classes=num_classes)
        self.valid_acc = torchmetrics.classification.Accuracy(task="multiclass", num_classes=num_classes)

    def training_step(self, batch, batch_idx):
        x, y = batch
        preds = self(x)
        self.train_acc(preds, y)
        # Logging the object directly
        self.log('train_acc', self.train_acc, on_step=True, on_epoch=False)

    def validation_step(self, batch, batch_idx):
        logits, y = batch
        self.valid_acc(logits, y)
        # Logging the object directly with epoch-level computation
        self.log('valid_acc', self.valid_acc, on_step=True, on_epoch=True)

Modular metrics in TorchMetrics are stateful objects that inherit from the Metric base class. They are designed for the typical deep learning workflow where data arrives in batches. Instead of requiring all predictions and targets to be available at once (like Scikit-learn), modular metrics use an online accumulation pattern:

1. self.add_state: Used in the __init__ method to define internal states that need to be accumulated. These states are automatically synchronized across multiple devices (GPUs/TPUs) in distributed environments.
2. update(preds, target): Accumulates the current batch of data into the global states.
3. compute(): Processes the accumulated states to return the final metric value.
4. forward(preds, target): Inherited from torch.nn.Module, this can be used to both get the metric for the current batch and accumulate the global state simultaneously.

This architecture allows for seamless scaling to multi-device training, as TorchMetrics handles the cross-device synchronization of the defined states.

import torch
from torch import tensor, Tensor
from torchmetrics import Metric

class Accuracy(Metric):
    def __init__(self):
        super().__init__()
        # Define states to be accumulated and synchronized
        self.add_state("correct", default=tensor(0), dist_reduce_fx="sum")
        self.add_state("total", default=tensor(0), dist_reduce_fx="sum")

    def update(self, preds: Tensor, target: Tensor) -> None:
        # Accumulate current batch into global states
        self.correct += torch.sum(preds == target)
        self.total += target.numel()

    def compute(self) -> Tensor:
        # Return final metric value from accumulated states
        return self.correct / self.total