面向分布式、可扩展 PyTorch 应用的机器学习指标库。

---

什么是 TorchMetrics • 实现自定义指标 • 内置指标 • 文档 • 社区 • 许可证

---

---

正在寻找 GPU 吗？

超过 34 万名开发者正在使用 Lightning Cloud，它专为 PyTorch 和 PyTorch Lightning 打造。

• GPU 低至每小时 0.19 美元。
• 集群：前沿级别的训练与推理集群。
• AI Studio（调试模式）：AI 助力调试、调优和高效训练的工作空间。
• AI Studio（部署模式）：AI 助力模型优化与部署的工作空间。
• Notebooks：持久化的 GPU 工作空间，AI 可帮助您编写代码并进行分析。
• 推理服务：将模型部署为推理 API。

安装

从 PyPI 简单安装：

pip install torchmetrics

conda install -c conda-forge torchmetrics

uv add torchmetrics

# 使用 git
pip install git+https://github.com/Lightning-AI/torchmetrics.git@release/stable

pip install https://github.com/Lightning-AI/torchmetrics/archive/refs/heads/release/stable.zip

pip install torchmetrics[audio]
pip install torchmetrics[image]
pip install torchmetrics[text]
pip install torchmetrics[all]  # 安装以上所有选项

pip install https://github.com/Lightning-AI/torchmetrics/archive/master.zip

什么是 TorchMetrics

TorchMetrics 是一个包含 100 多种 PyTorch 指标实现的库，并提供简单易用的 API 来创建自定义指标。它具有以下特点：

• 标准化接口，提升实验可重复性
• 减少样板代码
• 自动跨批次累积
• 针对分布式训练优化的指标
• 自动在多设备间同步

您可以将 TorchMetrics 与任何 PyTorch 模型一起使用，或结合 PyTorch Lightning 使用，以享受更多功能，例如：

• 模块级指标会自动放置到正确的设备上。
• 原生支持在 Lightning 中记录指标，进一步减少样板代码。

如何使用 TorchMetrics

模块级指标

基于模块的指标（参考文档）包含内部指标状态（类似于 PyTorch 模块的参数），能够自动完成跨设备的累积与同步！

• 自动跨多个批次累积
• 自动在多设备间同步
• 支持指标间的算术运算

无论是在 CPU、单 GPU 还是多 GPU 上，都可以运行！

对于单 GPU/CPU 场景：

import torch

# 导入我们的库
import torchmetrics

# 初始化指标
metric = torchmetrics.classification.Accuracy(task="multiclass", num_classes=5)

# 将指标移动到计算所需的设备上
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)

    # 当前批次的指标值
    acc = metric(preds, target)
    print(f"第 {i} 个批次的准确率: {acc}")

# 使用自定义累积方法计算所有批次的指标
acc = metric.compute()
print(f"所有数据上的准确率: {acc}")

在使用多 GPU 或多节点时，模块化指标的用法保持不变。

import 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"

    # 创建默认进程组
    dist.init_process_group("gloo", rank=rank, world_size=world_size)

    # 初始化模型
    metric = torchmetrics.classification.Accuracy(task="multiclass", num_classes=5)

    # 定义一个模型，并将你的指标附加到该模型上
    # 这样当对模型调用 .to(device) 时，指标的状态会被放置到正确的加速器上
    model = nn.Linear(10, 10)
    model.metric = metric
    model = model.to(rank)

    # 初始化 DDP
    model = DDP(model, device_ids=[rank])

    n_epochs = 5
    # 这里展示了对多个训练轮次的迭代
    for n in range(n_epochs):
        # 这里将被带有 DistributedSampler 的 DataLoader 替代
        n_batches = 10
        for i in range(n_batches):
            # 模拟一个分类问题
            preds = torch.randn(10, 5).softmax(dim=-1)
            target = torch.randint(5, (10,))

            # 当前批次的指标
            acc = metric(preds, target)
            if rank == 0:  # 只在 rank 0 打印
                print(f"第 {i} 个批次的准确率: {acc}")

        # 使用自定义累积方法计算所有批次和所有加速器上的指标
        # 准确率在两个加速器上是相同的
        acc = metric.compute()
        print(f"所有数据上的准确率: {acc}, 加速器 rank: {rank}")

        # 重置内部状态，使指标为新数据做好准备
        metric.reset()

    # 清理
    dist.destroy_process_group()


if __name__ == "__main__":
    world_size = 2  # 并行化的 GPU 数量
    mp.spawn(metric_ddp, args=(world_size,), nprocs=world_size, join=True)

实现你自己的模块化指标

实现你自己的指标非常简单，只需继承 torch.nn.Module 即可。只需继承 torchmetrics.Metric，并实现 update 和 compute 方法：

import torch
from torchmetrics import Metric


class MyAccuracy(Metric):
    def __init__(self):
        # 记得调用 super
        super().__init__()
        # 对于每个用于指标计算的内部状态，调用 `self.add_state`
        # dist_reduce_fx 表示用于从多个进程归约状态的函数
        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:
        # 提取预测的类别索引以计算准确率
        preds = preds.argmax(dim=-1)
        assert preds.shape == target.shape
        # 更新指标状态
        self.correct += torch.sum(preds == target)
        self.total += target.numel()

    def compute(self) -> torch.Tensor:
        # 计算最终结果
        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))

函数式指标

与 torch.nn 类似，大多数指标都有基于模块的版本和函数式版本。函数式版本是简单的 Python 函数，它们接受 torch.tensors 作为输入，并返回相应的指标作为 torch.tensor。

import torch

# 导入我们的库
import torchmetrics

# 模拟一个分类问题
preds = torch.randn(10, 5).softmax(dim=-1)
target = torch.randint(5, (10,))

acc = torchmetrics.functional.classification.multiclass_accuracy(
    preds, target, num_classes=5
)

涵盖的领域及示例指标

TorchMetrics 总共包含 100 多个指标，涵盖了以下领域：

• 音频
• 分类
• 检测
• 信息检索
• 图像
• 多模态（图像-文本-3D 谈话头像）
• 名义型
• 回归
• 分割
• 文本

每个领域可能需要一些额外的依赖项，可以通过 pip install torchmetrics[audio]、pip install torchmetrics['image'] 等命令安装。

其他功能

绘图

可视化指标对于理解机器学习算法的工作情况非常重要。Torchmetrics 内置了绘图支持（通过 pip install torchmetrics[visual] 安装依赖），几乎所有模块化指标都可通过 .plot 方法进行简单可视化！只需调用该方法即可获得任何指标的简单可视化效果！

import torch
from torchmetrics.classification import MulticlassAccuracy, MulticlassConfusionMatrix

num_classes = 3

# 这将生成两个分布，随着迭代次数增加，它们会越来越相似
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)

# 绘制单个值
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()

# 绘制多个值
values = []
for i in range(10):
    values.append(acc(preds(i), target(i)))
fig3, ax3 = acc.plot(values)

有关不同指标绘图的示例，请运行 此示例文件。

贡献吧！

Lightning + TorchMetrics 团队正在努力添加更多指标。但我们正在寻找像您一样出色的贡献者，提交新的指标并改进现有指标！

加入我们的 Discord 以获取成为贡献者的帮助！

社区

如需帮助或有任何问题，请加入我们在 Discord 上的庞大社区！

引用

我们很高兴能够延续开源软件的优良传统，多年来一直受到 Caffe、Theano、Keras、PyTorch、torchbearer、ignite、scikit-learn 和 fast.ai 等项目的启发。

如果您想引用本框架，可以使用 GitHub 内置的引用功能，根据此文件生成 BibTeX 或 APA 格式的引用（当然，前提是您非常喜欢它 😊）。

许可证

请遵守本仓库中列出的 Apache 2.0 许可协议。此外，Lightning 框架目前正处于专利申请中。

TorchMetrics 快速上手指南

TorchMetrics 是一个专为 PyTorch 设计的机器学习指标库，提供 100+ 种内置指标实现。它支持分布式训练、自动批次累积和多设备同步，能显著减少样板代码并提高实验可复现性。

环境准备

• 操作系统：Linux, macOS, Windows
• Python 版本：3.8+
• 核心依赖：
  • PyTorch >= 1.10.0
  • NumPy
  • Packaging
• 可选依赖：根据具体任务领域（如音频、图像、文本）可能需要额外安装对应扩展包。

安装步骤

基础安装

推荐使用 pip 进行安装：

pip install torchmetrics

国内加速安装

使用清华大学镜像源加速安装：

pip install torchmetrics -i https://pypi.tuna.tsinghua.edu.cn/simple

按需安装扩展包

针对特定领域的指标，可安装额外依赖：

# 音频指标
pip install torchmetrics[audio]

# 图像指标
pip install torchmetrics[image]

# 文本指标
pip install torchmetrics[text]

# 一次性安装所有扩展
pip install torchmetrics[all]

Conda 安装

conda install -c conda-forge torchmetrics

基本使用

TorchMetrics 提供两种主要使用方式：模块式（Module-based）和函数式（Functional）。推荐在训练循环中使用模块式，因为它能自动处理状态累积和设备同步。

1. 模块式用法（推荐用于训练）

模块式指标会自动累积多个 batch 的数据，并支持多 GPU 分布式训练。

import torch
import torchmetrics

# 初始化指标 (例如：多分类准确率)
metric = torchmetrics.classification.Accuracy(task="multiclass", num_classes=5)

# 将指标移动到计算设备 (CPU 或 CUDA)
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)

    # 更新当前 batch 的指标状态
    acc = metric(preds, target)
    print(f"Batch {i} Accuracy: {acc}")

# 计算所有累积数据的最终指标
final_acc = metric.compute()
print(f"Total Data Accuracy: {final_acc}")

# 如果需要开始新一轮评估，重置状态
metric.reset()

2. 函数式用法（适用于简单评估）

如果你只需要对单个 Tensor 进行一次性计算，可以使用函数式接口。

import torch
import torchmetrics

# 模拟数据
preds = torch.randn(10, 5).softmax(dim=-1)
target = torch.randint(5, (10,))

# 直接调用函数计算指标
acc = torchmetrics.functional.classification.multiclass_accuracy(
    preds, target, num_classes=5
)
print(f"Accuracy: {acc}")

3. 自定义指标

通过继承 torchmetrics.Metric 类，你可以轻松实现自定义指标。只需定义 update（更新状态）和 compute（计算结果）方法。

import torch
from torchmetrics import Metric

class MyAccuracy(Metric):
    def __init__(self):
        super().__init__()
        # 注册内部状态，dist_reduce_fx 指定分布式下的归约方式
        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:
        preds = preds.argmax(dim=-1)
        assert preds.shape == target.shape
        self.correct += torch.sum(preds == target)
        self.total += target.numel()

    def compute(self) -> torch.Tensor:
        return self.correct.float() / self.total

# 使用自定义指标
my_metric = MyAccuracy()
preds = torch.randn(10, 5).softmax(dim=-1)
target = torch.randint(5, (10,))

result = my_metric(preds, target)
print(f"Custom Accuracy: {result}")