Simple installation from PyPI
Other installationsInstall using conda
conda install -c conda-forge torchmetrics
Install using uv
Pip from source
# with git pip install git+https://github.com/Lightning-AI/torchmetrics.git@release/stable
Pip from archive
pip install https://github.com/Lightning-AI/torchmetrics/archive/refs/heads/release/stable.zip
Extra dependencies for specialized metrics:
pip install torchmetrics[audio] pip install torchmetrics[image] pip install torchmetrics[text] pip install torchmetrics[all] # install all of the above
Install latest developer version
pip install https://github.com/Lightning-AI/torchmetrics/archive/master.zip
TorchMetrics is a collection of 100+ PyTorch metrics implementations and an easy-to-use API to create custom metrics. It offers:
You can use TorchMetrics with any PyTorch model or with PyTorch Lightning to enjoy additional features such as:
The module-based metrics contain internal metric states (similar to the parameters of the PyTorch module) that automate accumulation and synchronization across devices!
This can be run on CPU, single GPU or multi-GPUs!
For the single GPU/CPU case:
import torch
# import our library
import torchmetrics
# initialize metric
metric = torchmetrics.classification.Accuracy(task="multiclass", num_classes=5)
# move the metric to device you want computations to take place
device = "cuda" if torch.cuda.is_available() else "cpu"
metric.to(device)
n_batches = 10
for i in range(n_batches):
# simulate a classification problem
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}")
Module metric usage remains the same when using multiple GPUs or multiple nodes.
Example using DDPimport os
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
from torch import nn
from torch.nn.parallel import DistributedDataParallel as DDP
import torchmetrics
def metric_ddp(rank, world_size):
os.environ["MASTER_ADDR"] = "localhost"
os.environ["MASTER_PORT"] = "12355"
# create default process group
dist.init_process_group("gloo", rank=rank, world_size=world_size)
# initialize model
metric = torchmetrics.classification.Accuracy(task="multiclass", num_classes=5)
# define a model and append your metric to it
# this allows metric states to be placed on correct accelerators when
# .to(device) is called on the model
model = nn.Linear(10, 10)
model.metric = metric
model = model.to(rank)
# initialize DDP
model = DDP(model, device_ids=[rank])
n_epochs = 5
# this shows iteration over multiple training epochs
for n in range(n_epochs):
# this will be replaced by a DataLoader with a DistributedSampler
n_batches = 10
for i in range(n_batches):
# simulate a classification problem
preds = torch.randn(10, 5).softmax(dim=-1)
target = torch.randint(5, (10,))
# metric on current batch
acc = metric(preds, target)
if rank == 0: # print only for rank 0
print(f"Accuracy on batch {i}: {acc}")
# metric on all batches and all accelerators using custom accumulation
# accuracy is same across both accelerators
acc = metric.compute()
print(f"Accuracy on all data: {acc}, accelerator rank: {rank}")
# Resetting internal state such that metric ready for new data
metric.reset()
# cleanup
dist.destroy_process_group()
if __name__ == "__main__":
world_size = 2 # number of gpus to parallelize over
mp.spawn(metric_ddp, args=(world_size,), nprocs=world_size, join=True) Implementing your own Module metric
Implementing your own metric is as easy as subclassing an torch.nn.Module. Simply, subclass torchmetrics.Metric and just implement the update and compute methods:
import torch
from torchmetrics import Metric
class MyAccuracy(Metric):
def __init__(self):
# remember to call super
super().__init__()
# call `self.add_state`for every internal state that is needed for the metrics computations
# dist_reduce_fx indicates the function that should be used to reduce
# state from multiple processes
self.add_state("correct", default=torch.tensor(0), dist_reduce_fx="sum")
self.add_state("total", default=torch.tensor(0), dist_reduce_fx="sum")
def update(self, preds: torch.Tensor, target: torch.Tensor) -> None:
# extract predicted class index for computing accuracy
preds = preds.argmax(dim=-1)
assert preds.shape == target.shape
# update metric states
self.correct += torch.sum(preds == target)
self.total += target.numel()
def compute(self) -> torch.Tensor:
# compute final result
return self.correct.float() / self.total
my_metric = MyAccuracy()
preds = torch.randn(10, 5).softmax(dim=-1)
target = torch.randint(5, (10,))
print(my_metric(preds, target))
Similar to torch.nn, most metrics have both a module-based and functional version. The functional versions are simple python functions that as input take torch.tensors and return the corresponding metric as a torch.tensor.
import torch
# import our library
import torchmetrics
# simulate a classification problem
preds = torch.randn(10, 5).softmax(dim=-1)
target = torch.randint(5, (10,))
acc = torchmetrics.functional.classification.multiclass_accuracy(
preds, target, num_classes=5
) Covered domains and example metrics
In total TorchMetrics contains 100+ metrics, which covers the following domains:
Each domain may require some additional dependencies which can be installed with pip install torchmetrics[audio], pip install torchmetrics['image'] etc.
Visualization of metrics can be important to help understand what is going on with your machine learning algorithms. Torchmetrics have built-in plotting support (install dependencies with pip install torchmetrics[visual]) for nearly all modular metrics through the .plot method. Simply call the method to get a simple visualization of any metric!
import torch
from torchmetrics.classification import MulticlassAccuracy, MulticlassConfusionMatrix
num_classes = 3
# this will generate two distributions that comes more similar as iterations increase
w = torch.randn(num_classes)
target = lambda it: torch.multinomial((it * w).softmax(dim=-1), 100, replacement=True)
preds = lambda it: torch.multinomial((it * w).softmax(dim=-1), 100, replacement=True)
acc = MulticlassAccuracy(num_classes=num_classes, average="micro")
acc_per_class = MulticlassAccuracy(num_classes=num_classes, average=None)
confmat = MulticlassConfusionMatrix(num_classes=num_classes)
# plot single value
for i in range(5):
acc_per_class.update(preds(i), target(i))
confmat.update(preds(i), target(i))
fig1, ax1 = acc_per_class.plot()
fig2, ax2 = confmat.plot()
# plot multiple values
values = []
for i in range(10):
values.append(acc(preds(i), target(i)))
fig3, ax3 = acc.plot(values)
For examples of plotting different metrics try running this example file.
The lightning + TorchMetrics team is hard at work adding even more metrics. But we're looking for incredible contributors like you to submit new metrics and improve existing ones!
Join our Discord to get help with becoming a contributor!
For help or questions, join our huge community on Discord!
We’re excited to continue the strong legacy of open source software and have been inspired over the years by Caffe, Theano, Keras, PyTorch, torchbearer, ignite, sklearn and fast.ai.
If you want to cite this framework feel free to use GitHub's built-in citation option to generate a bibtex or APA-Style citation based on this file (but only if you loved it 😊).
Please observe the Apache 2.0 license that is listed in this repository. In addition, the Lightning framework is Patent Pending.
RetroSearch is an open source project built by @garambo | Open a GitHub Issue
Search and Browse the WWW like it's 1997 | Search results from DuckDuckGo
HTML:
3.2
| Encoding:
UTF-8
| Version:
0.7.4
