torchmetrics实战：PyTorch多卡指标计算与分布式归约

做模型训练的人大概都有过这么一段经历训练脚本跑了一晚上最后打印出来的准确率漂亮得不像话换一台单卡机器重新推理一遍数字直接掉了七八个点。回头翻代码发现是自己写的一个累加器在多卡环境下每个进程各算各的均值最后又把这些均值平均了一次样本数不均衡的时候结果必然偏。这类问题跟模型结构、学习率、数据增强都没关系纯粹是指标计算这一层出了问题。torchmetrics就是在这个位置上干活的库它把准确率、F1、AUROC、BLEU、mAP 这些指标封装成带状态的nn.Module让累积、跨卡归约、跨 epoch 重置这几件琐事由框架接管。如果你正在用 PyTorch 写训练循环或者被多卡指标对不上、显存莫名爆掉、自己写的 IoU 慢得离谱这类问题折磨过下面的内容应该能省你不少时间。1. 先想清楚为什么指标代码不该自己手写1.1 三种常见写法的隐性成本大部分人算指标无非三条路手写累加器、调 sklearn、用 torchmetrics。表面看前两种更轻不用多装一个包但把维护成本折进去算账就不一样了。手写累加器的典型形态是一个字典加几个验证循环里correct (pred label).sum()epoch 结束correct / total。单卡单进程跑没问题一旦上 DDP每个 rank 只见到自己那份数据correct和total都是局部的最后要么只打印 rank0 的数字等于拿 1/N 的数据估整体要么各 rank 自己除完再 all_reduce 平均样本数不均衡时直接错。你还得处理total为 0、忽略 padding 位置、多分类的 macro/micro 区分等一堆分支。写的时候觉得二十行搞定改的时候发现每个指标都要重写一遍这套逻辑。sklearn 这条路在离线评估里很舒服classification_report一行出全表。但它在训练循环里有两个硬伤一是输入必须是 numpy 或 CPU 张量每步都得.cpu().numpy()GPU 到 CPU 的同步开销在几百步的小 batch 下非常可观二是它是全量计算你得把所有预测攒在内存里验证集上万条样本时显存/内存都得留一份做在线监控更是不可能。torchmetrics 的定位就在这两者之间像 sklearn 一样给出标准指标定义像手写累加器一样流式更新同时把分布式归约和设备迁移做进框架。维度手写累加器sklearn 全量计算torchmetrics单卡正确性容易写对正确正确多卡正确性需要自己实现 all_reduce且归约方式因指标而异需要先 gather 到主进程内建指定dist_reduce_fx即可设备手动搬运强制 CPU跟随模型自动.to(device)内存占用小只存中间量大存全量预测小存中间量如 TP/FP/FN指标覆盖面看你写多少经典 ML 指标齐全分类/回归/检索/文本/图像/音频/检测分割都有可微性无无部分指标支持反向如 SoftDice注意torchmetrics 不是 sklearn 的替代品。离线做一次性的深度分析、画 ROC 曲线、跑统计检验sklearn 和 scipy 依然更顺手。它的战场是训练与验证循环内部。1.2 设计取舍为什么做成有状态的 Module第一次看 torchmetrics 的 API 会觉得有点绕为什么不是accuracy(preds, labels)这样一个函数而要update一次、compute一次原因在于流式场景的硬约束。训练时数据是分 batch 流入的你没法把整个 epoch 的预测攒起来再算显存扛不住而且有些指标只需要中间统计量比如准确率只需要正确数和总数。所以指标对象必须能吃一口、记一笔最后一次性吐出结果。这就是update和compute的分工。那为什么继承nn.Module而不是普通类三个实际好处。第一nn.Module自带.to(device)指标内部的缓冲区会跟着走你在 GPU 上update传 GPU 张量不会报设备不匹配。第二它自带state_dict某些场景下你想把验证阶段的中间统计量存进 checkpoint 续跑直接可用。第三它能被nn.ModuleList或 Lightning 的self.log直接管理跟已有的工程体系无缝。代价也有每次.to(device)会遍历一遍内部缓冲区指标对象多的时候有一点点开销另外nn.Module的forward语义在这里被改写了见 2.2不熟悉的话容易踩。我个人判断一个指标库值不值得用看它敢不敢把归约方式暴露给用户。torchmetrics 暴露了dist_reduce_fx说明它认真处理了分布式这件事而不是在所有指标里硬编码一个mean。2. 把状态机吃透五个方法撑起整个库2.1 update 与 compute 的两段式流转一个指标对象的生命周期是这样的构造时通过add_state声明若干缓冲区比如tp、fp、fn、total初始值通常是 0 或空列表训练/验证循环里每步调update(preds, target)把这一批的统计量累加进缓冲区一个 epoch 结束调compute()从缓冲区算出最终标量或张量。关键在于累加的是什么。Accuracy 累加的是正确数和总数不是每批的准确率——这两者在 batch 大小不一致时结果完全不同。F1 累加的是 TP/FP/FN 三个计数不是每批的 F1 分数。AUROC 累加的是预测概率和标签因为它是排序指标必须看到全局分布才能算所以它的内部缓冲区会随样本数增长这是唯一需要注意内存的地方。写自定义指标时这条最容易被忽略你在update里存的必须是可加和的充分统计量不能存中间结果。我见过有人在update里直接算好 IoU 存进列表最后compute取平均小目标和大目标被同等加权结果跟标准实现差好几个点。2.2 forward 和 reset两个最容易被误用的方法forward(preds, target)等于update之后立刻compute返回当前累计的状态。它的用途是调试和单步评估不适合放在训练循环里——每步都算一次全量指标慢。reset()把缓冲区恢复到初始值。这个方法的调用时机是最常见的 bug 来源。标准节奏是验证循环开始前reset或者干脆每个 epoch 新建对象循环内update循环结束compute。如果你在 epoch 之间忘了reset验证指标会是所有 epoch 的累积值看起来越训越好其实是分母在膨胀。反过来在训练循环里每步都reset也是错的那样等于只统计最后一个 batch。提示用 Lightning 的self.log(val_acc, metric)时Lightning 会在 epoch 边界自动处理重置。自己写原生循环就必须手动管别指望框架。2.3 add_state 与分布式归约sum、mean、cat 该选哪个add_state是自定义指标的入口签名大致是add_state(name, default, dist_reduce_fxNone, persistentFalse)。dist_reduce_fx决定了多卡时这些缓冲区怎么合并dist_reduce_fx行为典型用途sum各卡缓冲区求和计数类TP/FP/FN、样本总数、平方误差和mean各卡求平均已经是均值的量如逐步 loss 的均值cat各卡拼接需要全局分布的指标AUROC 的概率、检索的 embeddingmin/max取极值最小/最大值跟踪None不归约各卡独立只在 rank0 使用的量选错的后果很直接。计数类指标用了mean各卡样本数不同时结果就是错的——这正好对应开头说的那个场景。排序类指标用了sum概率被加成一个巨大的数AUROC 直接失效。persistent控制这些缓冲区是否进入state_dict。默认False是对的指标中间状态通常不需要跟着模型存档存了反而让 checkpoint 变大续跑时还可能带进脏数据。2.4 sync 时机与 dist_sync_on_stepdist_sync_on_step是个布尔开关控制是否在每次update之后立刻做跨卡同步。默认False也就是只在compute时同步一次。打开它的场景很少你需要每一步都拿到全局一致的指标值来做某些在线决策比如早停的触发条件、动态调整采样权重。大多数情况下保持关闭因为每步一次 all_gather 在小 batch 训练里能吃掉明显的吞吐。同理指标对象内部还有sync_on_compute这类历史开关不同大版本之间有过调整。你在老项目里看到compute_on_step、sync_on_compute这些参数多半是 1.0 之前的写法升级时按当前版的文档对一遍别直接照抄旧博客的代码。3. 分类指标参数选错比不用还糟3.1 task、num_classes、average 三个参数的真实含义新版 torchmetrics 的分类指标统一用task参数区分场景binary、multiclass、multilabel。这个设计把过去散落在各处的num_classes、multilabel标志收敛成一处值得点赞。但三个参数之间的联动关系还是得搞清楚。taskmulticlass时必须给num_classes且它要和模型输出维度一致。给少了当预测类别索引超出范围时直接报错给多了混淆矩阵里多出几行几列全零指标数值不变但报表难看。taskmultilabel时num_classes表示标签维度preds应该是每个标签独立判断的 0/1阈值由你定或概率。average参数决定多分类怎么合成为一个数取值计算方式什么时候用micro全局汇总 TP/FP/FN 后再算类别不均衡、关心整体样本级表现macro每类单独算再取算术平均关心小类表现类别不均衡时的常用选择weighted每类单独算按支持度加权想接近 micro 但又保留类别信息none返回每类的值需要分类型诊断我见过的最典型的错误是数据里 95% 是负样本用micro的 F1 得到 0.94兴高采烈地汇报实际上模型把所有样本都预测成负类。macro会立刻暴露出正类的 F1 接近 0。这不是 torchmetrics 的问题但它把选择权交给你了选之前得知道自己要什么。3.2 logits、概率、标签输入格式对齐输入格式混乱是线上事故的高发区。粗略的规律是Accuracy、F1Score、Precision、Recall、ConfusionMatrix这类基于阈值的指标接受 logits 或概率内部会做argmax多分类或与阈值比较二分类。AUROC、AveragePrecision这类排序指标需要的是连续分数不是已经 argmax 过的类别索引。你把类别索引喂进去它照样能算出一个数但那个数没有意义。TopKC accuracy需要的是 logits 或概率矩阵k超过num_classes时会静默截断。二分类有个额外坑preds的形状可以是(N,)概率或(N, 1)但target的形状必须是(N,)的整型。形状对不上时有的版本报错有的版本广播出诡异的结果。我现在的习惯是统一在update前做一次preds.view(-1)和target.view(-1)省得后面排查。提示AUROC 这类指标不确定内部要不要做 sigmoid/softmax 时最稳的做法是自己显式转换后再传。显式转换不会有版本差异也让代码读起来一目了然。3.3 阈值类与排序类指标的取舍AUROC衡量的是排序质量与阈值无关适合评估模型区分正负的能力但不反映你在某个具体阈值下的业务表现。上线阶段更该看RecallAtFixedPrecision在精度达标的前提下召回多少或者PrecisionRecallCurve导出的曲线。多分类 AUROC 还有average的另一套语义macro是每类 one-vs-rest 后平均weighted按支持度加权跟 F1 的average不是同一个东西别搞混。3.4 StatScores 与 ConfusionMatrix一个指标拆成多个如果你的报表需要同时展示 precision、recall、F1、特异度最省算力的做法是只实例化一个StatScores或ConfusionMatrix从它的输出自己算派生指标。原因是这些指标共享同一套 TP/FP/FN 中间量分别实例化等于把同一份统计做四遍。当然更省事的是下一节要讲的MetricCollection的compute_groups它会自动识别共享状态的指标并合并计算。4. 分领域的指标清单与选型4.1 回归与相关性别拿 MSE 当唯一标准回归任务的常用指标在torchmetrics.regression下MeanSquaredError注意squaredFalse时输出 RMSE、MeanAbsoluteError、MeanSquaredLogError、R2Score、ExplainedVariance、PearsonCorrCoef、SpearmanCorrCoef、MeanAbsolutePercentageError。R2Score在多输出回归下有个multioutput参数raw_values、uniform_average、variance_weighted选错会得到差异很大的数字。MeanAbsolutePercentageError在目标值接近 0 时会爆炸做销量、流量这类稀疏目标时慎用改用MeanSquaredLogError或者对目标做 log 变换后再算。SpearmanCorrCoef内部要做排序大样本下比 Pearson 慢一个量级验证集几万条时注意频率别每步都算。4.2 检索与排序NDCG、MAP、HitRatetorchmetrics.retrieval覆盖了推荐和检索场景的常用指标RetrievalPrecision、RetrievalRecall、RetrievalMAP、RetrievalNormalizedDCG、RetrievalHitRate、RetrievalFallOut。这类指标的输入是indexes查询索引和preds相关性分数矩阵需要你显式告诉它哪些预测属于同一个 query。索引构造错了不会报错只会给出一个看似合理的错误数值。我一般的做法是先拿 10 条样本手动验算一遍确认索引映射正确后再跑全量。4.3 文本生成BLEU 与 ROUGE 的输入格式之争torchmetrics.text下有BLEUScore、ROUGEScore、chrFScore、EditDistance、BERTScore等。坑集中在输入格式上BLEU 接受 token 列表的列表ROUGE 接受字符串列表或 token 列表两者不一致。传错格式时不会崩只是分数接近 0你会以为是模型不行。另一个细节BLEU 的分词方式会显著影响分数英文按空格切问题不大中文必须先分词再用 token 输入直接传字符串进去得到的 n-gram 是错误的。做中文评估时我更倾向chrFScore字符级或者用BERTScore做语义相似度规避分词器差异。BERTScore会加载预训练模型首次调用有下载和加载开销务必在循环外实例化并.to(device)别在循环里新建。4.4 图像、音频、检测分割领域常用指标注意点图像重建PeakSignalNoiseRatio、StructuralSimilarityIndexMeasure输入需归一化到同一范围PSNR 的data_range要跟数据一致生成质量FrechetInceptionDistance、InceptionScore、KernelInceptionDistance需要足够样本数才稳定样本太少的 FID 没有参考价值音频SignalDistortionRatio、ScaleInvariantSignalNoiseRatio、PerceptualEvaluationSpeechQuality、ShortTimeObjectiveIntelligibility采样率必须与指标预期一致检测MeanAveragePrecision输入框格式xyxy/xywh和坐标是否归一化必须对齐分割IntersectionOverUnion、GeneralizedIntersectionOverUnion、Diceignore_index处理未标注像素MeanAveragePrecision底层依赖 coco 风格的评估逻辑构造时可以选择不同的后端速度快慢差异明显。检测任务的框格式是最容易出错的地方box_formatxyxy和xywh混用指标会安静地给你一个很低的 mAP你以为是模型没训好。5. 工程化接入MetricCollection 与训练循环5.1 MetricCollection 的 prefix/postfix 与 compute_groups指标一多代码里全是self.acc.update(...)、self.f1.update(...)这样的重复行。MetricCollection把这些收敛成一个字典式的容器import torchmetrics from torchmetrics import MetricCollection metrics MetricCollection({ acc: torchmetrics.Accuracy(taskmulticlass, num_classes10), f1_macro: torchmetrics.F1Score(taskmulticlass, num_classes10, averagemacro), f1_micro: torchmetrics.F1Score(taskmulticlass, num_classes10, averagemicro), auroc: torchmetrics.AUROC(taskmulticlass, num_classes10), }, prefixval_) metrics metrics.to(device)prefix和postfix在compute()返回的字典键上自动加前后缀省得自己拼字符串。compute_groupsTrue默认开启会自动把状态结构相同的指标归成一组共享一次update和一次跨卡归约。上面例子里的三个 F1/Accuracy 实际共享 TP/FP/FN归约只发生一次指标越多收益越明显。一个需要留意的副作用开启了compute_groups之后如果你手动对集合里的单个指标调用update可能触发重复计算或状态不一致。要么全走集合要么全走单体别混着来。5.2 原生 PyTorch 循环的标准接法import torch from torchmetrics import MetricCollection, Accuracy, F1Score def build_metrics(num_classes, device): return MetricCollection({ acc: Accuracy(taskmulticlass, num_classesnum_classes), f1: F1Score(taskmulticlass, num_classesnum_classes, averagemacro), }).to(device) train_metrics build_metrics(10, device) val_metrics build_metrics(10, device) for epoch in range(num_epochs): model.train() train_metrics.reset() for x, y in train_loader: x, y x.to(device), y.to(device) logits model(x) loss criterion(logits, y) loss.backward() optimizer.step() optimizer.zero_grad(set_to_noneTrue) train_metrics.update(logits, y) model.eval() val_metrics.reset() with torch.no_grad(): for x, y in val_loader: x, y x.to(device), y.to(device) logits model(x) val_metrics.update(logits, y) print({train: train_metrics.compute(), val: val_metrics.compute()})几个细节值得说。第一训练和验证用两个独立的指标对象不要复用否则验证指标会混进训练数据。第二reset()放在循环开头而不是compute()之后这样即使中间抛异常下一轮也不会带着脏状态。第三update传的是 logitsAccuracy 内部会做 argmax不用自己处理。第四compute()的结果是张量打日志前记得.item()或者交给日志库处理别把 GPU 张量传进不需要它的地方。5.3 与 Lightning 的配合Lightning 用户更省心直接把指标对象挂到self上在training_step/validation_step里update在on_train_epoch_end/on_validation_epoch_end里compute并self.log。Lightning 会在合适的时机重置也会自动处理 DDP 下的同步。需要注意的是self.log的on_epoch、on_step、sync_dist三个参数组合。默认on_stepTrue, on_epochFalse会导致日志里塞满每步的值验证阶段建议on_stepFalse, on_epochTrue。5.4 性能与显存什么时候该 compute_on_cputorchmetrics 的指标默认在 GPU 上累积。大多数计数类指标只有几个标量显存开销可以忽略。但排序类AUROC、AP和检索类指标会把概率或分数存进缓冲区样本量大时会占可观显存极端情况下能和模型激活值抢资源。处理方法有两个一是降低update频率每 N 步更新一次二是对这类指标调用compute_on_cpu()让缓冲区落在内存里。代价是每次update都要做一次 GPU→CPU 拷贝小 batch 下这个拷贝的同步成本不低。我的经验是验证集规模在万级以内GPU 累积完全没问题跑到十万级、且有多个排序类指标同时开就该考虑搬 CPU 或者抽样评估了。具体阈值跟显卡和 batch 大小强相关实测一遍比我给数字靠谱。6. 自定义指标从骨架到分布式安全6.1 一个可用的骨架import torch from torchmetrics import Metric class WeightedAccuracy(Metric): def __init__(self, class_weights, dist_sync_on_stepFalse): super().__init__(dist_sync_on_stepdist_sync_on_step) self.class_weights torch.tensor(class_weights, dtypetorch.float) # 分子、分母都是可加和的计数用 sum 归约 self.add_state(weighted_correct, defaulttorch.tensor(0.0), dist_reduce_fxsum) self.add_state(weighted_total, defaulttorch.tensor(0.0), dist_reduce_fxsum) def update(self, preds, target): preds preds.argmax(dim-1) if preds.ndim 1 else preds w self.class_weights.to(target.device)[target] self.weighted_correct (w * (preds target).float()).sum() self.weighted_total w.sum() def compute(self): return self.weighted_correct / self.weighted_total这个骨架里三个关键点add_state声明的缓冲区会自动参与.to(device)和跨卡归约update里只累加可加和的量compute里只做最终除法不改缓冲区。6.2 五个必须检查的点第一__init__里不要创建需要跟随设备的张量而不注册成 state。放在普通属性里的张量不会跟着.to(device)走后面设备不匹配就报错。上面例子里的class_weights在update里手动.to(target.device)是权宜写法更规范的做法是注册成 state 或者用register_buffer。第二compute里不要原地修改 state。需要归一化就先.clone()。torchmetrics 某些版本对compute的副作用检查不严你在本地跑对了多卡下可能因为第二次调用compute结果不同而报错。第三dist_reduce_fx要跟统计量语义匹配。计数用sum已经除过的比例用mean需要全局分布的用cat。第四注意空状态。第一个 batch 之前就调compute分母是 0结果是 nan。加一个极小值保护或者在使用侧保证顺序。第五update内部不要做跨卡通信。有人图省事在update里写all_reduce结果每步都同步训练慢到怀疑人生。同步交给框架在compute时统一做。7. 排查速查表我踩过的坑7.1 现象到原因现象最可能的原因处理方式多卡指标比单卡系统性偏低计数类 state 用了mean归约改成dist_reduce_fxsum验证指标每轮递增、收敛后还在涨epoch 之间没reset循环开头 reset 或每轮新建AUROC 恒为 0.5 左右传入了 argmax 后的类别索引改传概率或 logits二分类 F1 异常低target形状是(N,1)或类型是 float统一view(-1).long()训练越跑越慢、显存缓慢增长排序类指标把全部预测存进 GPU 缓冲区compute_on_cpu()或降低 update 频率Expected all tensors on same device自定义指标里用了未注册的普通属性张量改用add_state或register_buffer指标集合里某个指标数值不动混用集合与单体 update 导致状态分裂统一走MetricCollection中文 BLEU 接近 0直接传字符串n-gram 切分错误先分词传 token 列表7.2 几条长期有效的习惯写完训练循环先做一次单卡 vs 多卡的数值对齐验证同一份小数据集几十条单卡跑一遍、双卡跑一遍指标差值应该在浮点误差范围内。这一步花五分钟能省掉后面几天的困惑。指标对象的构造统一收进一个工厂函数把num_classes、task、average、device这些散落的参数集中管理。模型改了输出维度只改一处。验证阶段的指标只算必要的。全量报表每类的 P/R/F1、混淆矩阵、PR 曲线放在独立的评估脚本里跑别塞进训练循环那样既慢又容易被中途打断。写完一个自定义指标拿 sklearn 的对应实现对一遍。用小样本100 条左右、随机预测两个实现的差值在 1e-6 以内才算过。这一步能抓出 90% 的归约和形状错误。我在实际项目里的体会是torchmetrics 的价值不在于它替你省了多少行代码而在于它把指标计算从一件每人都重写一遍的事变成了一个有明确契约的接口update吃什么、compute吐什么、跨卡怎么合并都有章可循。真正花时间的从来不是调包而是搞清楚自己到底要衡量什么——average选 macro 还是 micro评估排序质量还是阈值表现这些问题框架替不了你。把这些想明白之后再去看那些参数就不会觉得是在背 API 了。