pytorch-metric-learning 测试器 5 步跑通

pytorch-metric-learning 测试器 5 步跑通

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

pytorch-metric-learning 的 Testers 负责嵌入评估：批量算嵌入、最近邻检索、出 precision_at_1 与 NMI 等指标。从定位到训练闭环，带你跑通。

🔍 先搞清楚角色：Testers 做什么、不做什么

Testers 是个纯评估组件，把它想成"自动阅卷系统"。训练交给 trainers，它不碰任何梯度更新，只干三件事：

1. 拿模型和数据集，批量算出每个样本的嵌入向量；
2. 在嵌入空间里做最近邻（k-NN）检索；
3. 输出 precision_at_1、NMI 这类精度指标。

它有一个硬依赖：faiss。底层 k-NN 检索（FaissKNN）和 k-means 聚类（FaissKMeans，算 NMI/AMI 用）都封装在 faiss 里。装一行：

conda install faiss-cpu # 有卡就换 faiss-gpu

不装 faiss，默认的AccuracyCalculator跑不起来。

三步拿到第一份评估结果

第一步，建实例。挑一个具体测试器——日常验证基本用GlobalEmbeddingSpaceTester。它继承自BaseTester，那 15 个参数全都能接。

第二步，调 test()。核心就这一行：

from pytorch_metric_learning import testers dataset_dict = {"train": train_ds, "val": val_ds} acc = testers.GlobalEmbeddingSpaceTester().test(dataset_dict, epoch, model)

dataset_dict是"名字 → 数据集"的映射，epoch只进日志。模型有俩传法：单model，或trunk+embedder分开传。

后者里embedder缺省会退化成torch.nn.Identity()（无操作），直接拿 trunk 输出当嵌入。

第三步，读 all_accuracies。结果是嵌套字典all_accuracies[split][metric]，指标名带level0后缀（指第 0 层标签层级，单标签数据可忽略）。想自定义收尾，就挂一个end_of_testing_hook——它在所有 split 跑完后只触发一次：

# all_accuracies 长这样：[split][metric] {"val": {"precision_at_1_level0": 0.38, "NMI_level0": 0.44}} def hook(tester): print(tester.all_accuracies)

完整调用链在src/pytorch_metric_learning/testers/base_tester.py，想看细节可直接翻。

三个内置测试器怎么选

• Global：把整个嵌入空间当题库，查谁都不受限，日常验证基本用它。
• WithSameParent：检索收窄到"同一父类"内部，比如 {cat, dog} 只在 animal 里找、{car, truck} 只在 vehicle 里找，最后各父类指标取平均。
• TwoStream：专给双流损失用，数据集得吐三元组，且 query 必须等于唯一的 reference。

BaseTester 参数速查

15 个参数按用途分四组。

数据输入组

计算控制组

标签处理组

输出与可视化组

补两点：label_hierarchy_level取"all"或列表时会额外产出AVERAGE_前缀的汇总指标；set_min_label_to_zero开字符串标签（'dog'/'cat' 会被映成整数）时，必须同时传dataset_labels。

splits_to_eval 进阶组合

默认splits_to_eval=None，等价于每个 split 拿自己当 query、也当 reference：[(k, [k]) for k in dataset_dict]。

query 是"考题"，reference 是"题库"。想让两者不一样，就传元组列表(query, [refs])：

# 考题来自 test，题库来自 train tester.test(ds, epoch, model, splits_to_eval=[("test", ["train"])]) # 考题来自 a，题库 = a + train 合并 tester.test(ds, epoch, model, splits_to_eval=[("a", ["a", "train"])])

几条规则：query split 不能重复出现；要算嵌入的 split 会被去重、每个只前向一次；若 query 本身在 reference 里，会被挪到 reference 最前面。

原因是get_accuracy()要求ref_includes_query=True时，reference 前段必须与 query 逐元素一致。

一个容易忽略的点：默认自查询会让指标略微虚高，因为样本总能检索到自己。想测泛化，就让 query 来自测试集、reference 来自训练集。

自定义指标与训练闭环

扩展指标。测试器不自己写精度数学，都委托给AccuracyCalculator。默认指标：NMI/AMI（用FaissKMeans对 query 嵌入聚类，簇数等于唯一标签数）、precision_at_1、r_precision、mean_average_precision_at_r等。

某些 query 标签若没出现在 reference 里，对应样本会被排除出 k-NN 类指标——它们拿不到非零的最近邻精度，硬算 0 分反而失真。

想加自定义指标：写一个AccuracyCalculator子类，方法以calculate_开头，声明requires_knn()/requires_clustering()，再传给测试器；用include/exclude可只算部分指标提速。自定义label_comparison_fn（非torch.eq）时，要排除聚类类指标 NMI/AMI。

接线训练闭环。把end_of_testing_hook接到 logging_presets，就得到"每轮验证 + 存最佳 + 早停"：

import pytorch_metric_learning.utils.logging_presets as LP rk, _, _ = LP.get_record_keeper("logs", "tb") hooks = LP.get_hook_container(rk) tester = testers.GlobalEmbeddingSpaceTester(end_of_testing_hook=hooks.end_of_testing_hook) eoe = hooks.end_of_epoch_hook(tester, {"val": val_ds}, "saved_models", test_interval=1, patience=5)

primary_metric决定哪个指标挑最佳 checkpoint，默认mean_average_precision_at_r（可选r_precision、precision_at_1、NMI）；patience控早停。结果同时落到 csv、sqlite、tensorboard，模型和优化器存进指定文件夹。

⚡ 避坑与性能优化清单

1. GPU 没被用上→ faiss-gpu 仅k≤1024（CUDA<9.5）或k≤2048（CUDA≥9.5）才走 GPU，否则回退 CPU 并警告 → 减小 k，或接受 CPU 回退。
2. 装了 faiss-cpu 恒走 CPU→ 装的是 CPU 版，天然没有 GPU 路径 → 有卡就换faiss-gpu。
3. 字符串标签对不上→LabelMapper缺dataset_labels无法重映射 →set_min_label_to_zero=True并传dataset_labels。
4. 大数据集检索慢→ 默认 k 等于 reference 样本总数 → 设k=1000或k="max_bin_count"。
5. r_precision / MAP@r 数值失真→k < max(bincount(reference_labels))→ 把 k 提到最大类频数以上。
6. test() 后取 embeddings_and_labels 报错→ 该属性在test()末尾被del→ 只在end_of_testing_hook里访问。
7. 多层级想要汇总分→ 单层只出单层级指标 →label_hierarchy_level设"all"或列表，产出AVERAGE_前缀指标。

把 Testers 想成"算嵌入 → 查最近邻 → 出分"的阅卷机，配上 AccuracyCalculator 和 logging_presets，就能在训练里自动挑出最强的 checkpoint；想深挖行为约定，tests/testers/下的单元测试是最好的延伸阅读。

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