pytorch-metric-learning 测试器 5 步跑通【免费下载链接】ocpFAIR Chemistrys library of machine learning methods for chemistry项目地址: https://gitcode.com/GitHub_Trending/oc/ocppytorch-metric-learning 的 Testers 负责嵌入评估批量算嵌入、最近邻检索、出 precision_at_1 与 NMI 等指标。从定位到训练闭环带你跑通。 先搞清楚角色Testers 做什么、不做什么Testers 是个纯评估组件把它想成自动阅卷系统。训练交给 trainers它不碰任何梯度更新只干三件事拿模型和数据集批量算出每个样本的嵌入向量在嵌入空间里做最近邻k-NN检索输出 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或trunkembedder分开传。后者里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想看细节可直接翻。三个内置测试器怎么选测试器输入形态检索范围适用场景限制GlobalEmbeddingSpaceTester单个嵌入向量全嵌入空间通用单流度量学习默认首选无WithSameParentLabelTester二维标签细粒度父标签仅同父标签的兄弟样本层级标签的细粒度评估标签须二维只能评L 维度-1的层级GlobalTwoStreamEmbeddingSpaceTester(anchor, positive, label)三元组anchor 查 posneg配TwoStreamMetricLoss评估仅支持自查询自参考Global把整个嵌入空间当题库查谁都不受限日常验证基本用它。WithSameParent检索收窄到同一父类内部比如 {cat, dog} 只在 animal 里找、{car, truck} 只在 vehicle 里找最后各父类指标取平均。TwoStream专给双流损失用数据集得吐三元组且 query 必须等于唯一的 reference。BaseTester 参数速查15 个参数按用途分四组。数据输入组参数默认值一句话说明何时需要改batch_size32算嵌入时每批样本数显存紧就调小、求快就调大dataloader_num_workers2数据加载进程数加载成瓶颈时data_deviceNone数据放哪None 自动选 GPU/CPU想强制指定设备dtypeNone数据输出转成什么类型如 float16低精度推理时data_and_label_getterNone把__getitem__输出拆成(data, label)的函数返回格式不是(data, label)时pcaNone嵌入 PCA 降维到的维数高维可视化/提速检索keyname 会带_pcaN计算控制组参数默认值一句话说明何时需要改normalize_embeddingsTrue最近邻前把嵌入 L2 归一化用余弦/欧氏距离检索时保持默认use_trunk_outputFalse直接用 trunk 输出、跳过 embedderembedder 是 no-op 或想单看 trunkaccuracy_calculatorNone用默认算精度的AccuracyCalculator实例自定义指标 / 调 k 加速标签处理组参数默认值一句话说明何时需要改label_hierarchy_level0用哪层标签支持 0 / all / 列表多层级标签时dataset_labelsNone数据集的标签数组set_min_label_to_zeroTrue时必填set_min_label_to_zeroFalse把标签映射成排名0,1,2…用字符串标签时输出与可视化组参数默认值一句话说明何时需要改end_of_testing_hookNone全部 split 跑完后触发一次的钩子想自定义收尾存日志/画图visualizerNone带fit_transform的对象如 UMAP降到 2D想看嵌入聚类结构visualizer_hookNone拿到降维结果后执行的回调配合 visualizer 存散点图补两点label_hierarchy_level取all或列表时会额外产出AVERAGE_前缀的汇总指标set_min_label_to_zero开字符串标签dog/cat 会被映成整数时必须同时传dataset_labels。splits_to_eval 进阶组合默认splits_to_evalNone等价于每个 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_queryTrue时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_hookhooks.end_of_testing_hook) eoe hooks.end_of_epoch_hook(tester, {val: val_ds}, saved_models, test_interval1, patience5)primary_metric决定哪个指标挑最佳 checkpoint默认mean_average_precision_at_r可选r_precision、precision_at_1、NMIpatience控早停。结果同时落到 csv、sqlite、tensorboard模型和优化器存进指定文件夹。⚡ 避坑与性能优化清单GPU 没被用上→ faiss-gpu 仅k≤1024CUDA9.5或k≤2048CUDA≥9.5才走 GPU否则回退 CPU 并警告 → 减小 k或接受 CPU 回退。装了 faiss-cpu 恒走 CPU→ 装的是 CPU 版天然没有 GPU 路径 → 有卡就换faiss-gpu。字符串标签对不上→LabelMapper缺dataset_labels无法重映射 →set_min_label_to_zeroTrue并传dataset_labels。大数据集检索慢→ 默认 k 等于 reference 样本总数 → 设k1000或kmax_bin_count。r_precision / MAPr 数值失真→k max(bincount(reference_labels))→ 把 k 提到最大类频数以上。test() 后取 embeddings_and_labels 报错→ 该属性在test()末尾被del→ 只在end_of_testing_hook里访问。多层级想要汇总分→ 单层只出单层级指标 →label_hierarchy_level设all或列表产出AVERAGE_前缀指标。把 Testers 想成算嵌入 → 查最近邻 → 出分的阅卷机配上 AccuracyCalculator 和 logging_presets就能在训练里自动挑出最强的 checkpoint想深挖行为约定tests/testers/下的单元测试是最好的延伸阅读。【免费下载链接】ocpFAIR Chemistrys library of machine learning methods for chemistry项目地址: https://gitcode.com/GitHub_Trending/oc/ocp创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考