torchmetrics

TorchMetrics项目介绍

简介

TorchMetrics是一个专为PyTorch应用设计的、高效且可扩展的机器学习指标库。它包括超过100种的内置指标，同时提供了一套易于使用的API，方便用户创建自定义指标。其主要特点包括：

• 提供标准化接口，提高了模型结果的可重复性
• 减少样板代码，简化开发流程
• 自动累积批处理数据，优化分布式训练
• 在多个设备间自动同步数据

用户可以在任何PyTorch模型中使用TorchMetrics，也可以结合PyTorch Lightning一起使用，以便享受额外的功能，如自动将模块化指标放置在正确的设备上并支持记录指标。

安装

TorchMetrics的安装非常简单。可以通过以下几种方法进行安装：

• 使用PyPI安装：

  pip install torchmetrics
  

• 使用Conda进行安装：

  conda install -c conda-forge torchmetrics
  

• 从源码安装：

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

• 安装额外依赖以支持不同领域的指标（如音频、图像、文本等）：

  pip install torchmetrics[audio]
  pip install torchmetrics[image]
  pip install torchmetrics[text]
  pip install torchmetrics[all]  # 下载安装所有依赖
  

使用指南

模块化指标

模块化指标利用内部的指标状态（类似于PyTorch模块的参数），可以在多个设备中自动进行累积和同步。以下是示例代码，展示了如何在单个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}")

自定义指标

用户可以通过继承torchmetrics.Metric来实现自己的指标。只需实现update和compute方法即可：

import torch
from torchmetrics import Metric

class MyAccuracy(Metric):
    def __init__(self):
        super().__init__()
        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))

功能性指标

除了模块化指标，每个指标也有其功能性版本，这些版本是简单的Python函数，输入torch.tensors，输出相应的指标结果。

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个指标，涵盖了以下领域：

• 音频
• 分类
• 检测
• 信息检索
• 图像
• 多模态（图文）
• 名义数据
• 回归
• 分割
• 文本

附加功能

绘图

TorchMetrics支持大部分指标的可视化。借助这个功能，用户可以快速了解他们的机器学习算法的工作情况。

贡献

TorchMetrics团队正在努力添加更多的指标，也需要来自社区的贡献。如果有兴趣贡献新的指标或者改进现有的指标，欢迎加入我们的Discord社区。