pytorch-metric-learning Testers 完整指南：从选型到训练闭环，一次讲透度量学习模型评估

pytorch-metric-learning Testers 完整指南：从选型到训练闭环，一次讲透度量学习模型评估

【免费下载链接】ocpFAIR Chemistry's library of machine learning methods for chemistry项目地址: https://gitcode.com/GitHub_Trending/oc/ocp

pytorch-metric-learning 的 Testers（评估器）是度量学习链路里的"考官"：拿到训练好的模型（trunk + embedder）和数据集划分后，它为每个样本批量算出嵌入向量（可以理解为"样本指纹"），再基于最近邻检索给模型打出一组精度分数。本文按"先选型 → 走链路 → 读结果 → 接闭环"的顺序展开，覆盖 Testers 的完整用法、指标输出解读，以及如何用 logging_presets 实现"训练—验证—保存最佳模型"的闭环。

速览：评估器卡在"数据 → 嵌入 → 检索 → 指标"的哪一环

Testers 不参与梯度更新，是纯评估组件。它的内部流程固定为四步：

1. 嵌入生成：把dataset_dict里每个 split 的样本过一遍模型（eval()模式 + 无梯度），得到高维嵌入；
2. 检索：默认用 faiss 后端（FaissKNN）在 reference 集合中为每条 query 做 k-NN；
3. 指标计算：把 query/reference 的嵌入与标签交给AccuracyCalculator，产出precision_at_1、r_precision、NMI 等指标；
4. 输出：结果写入tester.all_accuracies（嵌套字典：split → 指标名 → 数值），并可触发可视化与end_of_testing_hook。

一句话总结：你提供模型和数据，它交回一张"分数表"。下这张图是模型评估对比图的常见画法（一轴是误差指标、一轴是推理开销），理解 Testers 输出时可以套用同样的视角：一张表里同时回答"准不准"和"贵不贵"。

调用形态非常统一，所有测试器都是这么用的：

from pytorch_metric_learning import testers t = testers.GlobalEmbeddingSpaceTester() dataset_dict = {"train": train_ds, "val": val_ds} all_accuracies = t.test(dataset_dict, epoch, trunk, embedder) # embedder 可省略

先选型：我该用哪个 Tester

结论先行：90% 的场景用GlobalEmbeddingSpaceTester；只有当标签带层级、或你的损失函数是TwoStreamMetricLoss时，才需要另外两个。三者的边界如下（源码位于 src/pytorch_metric_learning/testers/）：

两个容易踩的选型细节：

• TwoStream 版内部把每个 batch 的 anchor 嵌入拼在前半段当 query、posneg 嵌入拼在后半段当 reference，且ref_includes_query恒为False（两边不是同一批样本）。它会在set_reference_and_query()里显式断言 query 与 reference 是同一个 split，传自定义组合会直接失败。
• WithSameParentLabel 版对每个唯一父标签筛出子集、在子集内分别算指标，最后跨父标签取平均（np.mean）。它衡量的是"细粒度类在粗粒度类内部的可分性"，与全局视角回答的不是同一个问题。

从配置到产出的完整链路

第一步：嵌入生成——指纹得先"可比"

BaseTester（src/pytorch_metric_learning/testers/base_tester.py）的__init__参数决定了这一步的行为，最常用的几个：

两点补充：

• 不传embedder_model时，源码会把它替换成torch.nn.Identity()，即"无操作"，直接消费 trunk 输出；
• 嵌入生成逻辑封装在get_all_embeddings()里，可脱离test()单独调用——想自己导出嵌入做 t-SNE/UMAP，用它就够了。

第二步：query 与 reference 的组合技巧

默认行为是"自查询"：dataset_dict里每个 split 既当 query 又当 reference。更灵活的组合通过test()的splits_to_eval参数指定，它是(query_split, [reference_splits])元组的列表：

# 默认等价于： [("dataset_a", ["dataset_a"]), ("train", ["train"])] # query 用测试集，reference 用训练集（评估泛化）： [("test", ["train"])] # reference 由多个 split 合并： [("test", ["test", "train"])]

源码在get_splits_to_compute_embeddings()中处理三件事：query split 不允许重复出现；query 与 reference 的并集才计算嵌入（同一 split 只跑一次前向）；若 query 包含在 reference 列表中，会被移到 reference 最前面——因为AccuracyCalculator.get_accuracy()要求ref_includes_query=True时，reference 的前len(query)个元素必须与 query 逐元素相等（有torch.allclose显式校验）。

经验建议：关心泛化能力时，让 query 来自测试集、reference 来自训练集，避免自查询带来的"样本总能检索到自身"导致的指标虚高。

第三步：kNN 检索——faiss 是引擎

检索由 src/pytorch_metric_learning/utils/inference.py 中的FaissKNN（k-NN）与FaissKMeans（NMI/AMI 用的聚类）完成，也是AccuracyCalculator的默认后端，因此faiss 是硬依赖。

关键参数是k：

• 默认k=None，等价于 k = reference 嵌入总数——数据集一大，检索就会很慢；
• 推荐显式传k=1000或k="max_bin_count"（每个类别的最大样本数），显著提速。

装faiss-gpu时，只有k <= 1024（CUDA < 9.5）或k <= 2048（CUDA ≥ 9.5）才走 GPU，否则回退 CPU 并告警——小数据集上若发现检索没跑在 GPU 上，多半是这个原因。

第四步：指标计算——考官判分

检索结果交给 src/pytorch_metric_learning/utils/accuracy_calculator.py 中的AccuracyCalculator。默认指标速览：

一个细节：若某些 query 标签根本没出现在 reference 里，这些样本会被排除出 k-NN 类指标——它们注定 0 分，不能反映嵌入质量，排除后更公平。

读懂输出：keyname 规则、指标值与"自查询虚高"

test()返回（并缓存在tester.all_accuracies上）的结构是两层字典：

{"val": {"precision_at_1_level0": 0.42, "NMI_level0": 0.55, ...}}

指标名遵循固定拼装规则（accuracies_keyname()生成）：

• _levelL后缀：L 是label_hierarchy_level参数选中的标签层级。单标签数据保持默认 0 即可，直接忽略后缀；
• AVERAGE_前缀：当label_hierarchy_level="all"或传入层级列表时，会对各层级额外汇总出AVERAGE_precision_at_1这类平均指标；
• _pcaN后缀：设置了pca后自动追加，标明该组指标基于降维后的嵌入。

关于"自查询虚高"：默认自查询模式下，每个样本的最近邻大概率是自己，precision_at_1天然偏高。这不是 bug，而是口径问题——对比不同模型时保证口径一致即可，单独汇报泛化数字时则建议用"test 查 train"的组合（见上文splits_to_eval示例）。

另外，epoch参数只用于日志与历史记录（写入all_accuracies[split]["epoch"]），不影响任何计算。

接入训练闭环：logging_presets + 早停 + 保存最佳模型

评估器的end_of_testing_hook最实用的去处，是接上 docs/logging_presets.md 里的logging_presets，把"每轮验证 + 挑最佳模型 + 早停"自动化：

import pytorch_metric_learning.utils.logging_presets as LP rk, _, _ = LP.get_record_keeper("logs", "tensorboard_logs") hooks = LP.get_hook_container(rk) dataset_dict = {"val": val_dataset} tester = testers.GlobalEmbeddingSpaceTester(end_of_testing_hook=hooks.end_of_testing_hook) end_of_epoch_hook = hooks.end_of_epoch_hook(tester, dataset_dict, "saved_models") trainer = trainers.MetricLossOnly( models, optimizers, batch_size, loss_funcs, mining_funcs, train_dataset, sampler=sampler, end_of_iteration_hook=hooks.end_of_iteration_hook, end_of_epoch_hook=end_of_epoch_hook, ) trainer.train(num_epochs=num_epochs)

要点：

• HookContainer的primary_metric决定用哪个指标挑最佳 checkpoint，默认mean_average_precision_at_r，可选r_precision、precision_at_1、NMI；
• end_of_epoch_hook返回真正的钩子函数，验证集键名默认"val"，还支持test_interval（隔几轮验证）与patience（早停耐心值）；
• 历史记录同时落 csv、sqlite 与 tensorboard，之后可用hooks.get_accuracy_history(tester, "val", ...)回放任意指标的历史曲线。

避坑与提速清单

• faiss 必装：FaissKNN/FaissKMeans/PCA 都压在它身上，用 conda 安装faiss-cpu或faiss-gpu；
• 大集提速：显式设accuracy_calculator=AccuracyCalculator(k=1000)或k="max_bin_count"；注意r_precision/mean_average_precision_at_r在 k 偏小时会失真；
• 字符串标签：必须set_min_label_to_zero=True并同时传dataset_labels，LabelMapper才会做重映射（此时标签按集合内排名变成连续整数）；
• 双流测试器的硬限制：splits_to_eval只支持默认值，query 与 reference 必须是同一 split，别尝试传自定义组合；
• 临时数据别出 hook：embeddings_and_labels在test()末尾被删除释放内存，需要在里面读嵌入就放在end_of_testing_hook里做；
• 钩子只调用一次：end_of_testing_hook在所有 split 评估完之后触发一次，不是每个 split 一次；
• 只算部分指标：AccuracyCalculator(include=("precision_at_1",))可显著加速评估。

下一步

• 跑一遍 examples/notebooks 下的示例 notebook（如 TripletMarginLossMNIST），观察每个 epoch 的验证指标如何进入all_accuracies并驱动最佳模型保存；
• 细读 tests/testers/ 下的三个测试文件：splits_to_eval三种组合对precision_at_1数值的影响、PCA 分支的端到端行为、双流测试器的断言，都在断言里写得明明白白；
• 深入指标数学见 docs/accuracy_calculation.md，参数总览见 docs/testers.md。

【免费下载链接】ocpFAIR Chemistry's library of machine learning methods for chemistry项目地址: https://gitcode.com/GitHub_Trending/oc/ocp