allRank:基于PyTorch的学习排序框架
关于
allRank是一个基于PyTorch的框架,用于训练神经网络学习排序(LTR)模型,具有以下特点:
- 常见的逐点式、成对式和列表式损失函数
- 全连接和Transformer式评分函数
- 常用评估指标,如归一化折扣累积增益(NDCG)和平均倒数排名(MRR)
- 用于模拟点击率数据实验的点击模型
动机
allRank提供了一种简单灵活的方式来实验各种LTR神经网络模型和损失函数。 添加自定义损失函数很容易,并且可以方便地配置模型和训练过程。 我们希望allRank能够促进神经网络LTR的研究及其工业应用。
特点
已实现的损失函数:
- ListNet(适用于二元和分级相关性)
- ListMLE
- RankNet
- 序数损失
- LambdaRank
- LambdaLoss
- ApproxNDCG
- RMSE
- NeuralNDCG(介绍于 https://arxiv.org/pdf/2102.07831)
入门指南
为了帮助您入门,我们提供了一个run_example.sh
脚本,该脚本生成libsvm格式的虚拟排序数据,并使用提供的示例config.json
配置文件在该数据上训练Transformer模型。运行脚本后,可以在dummy_data
目录中找到虚拟数据,在test_run
目录中找到实验结果。运行示例需要Docker。
获取正确的架构版本(GPU vs 仅CPU)
由于torch二进制文件在GPU和CPU上不同,且GPU版本不适用于CPU,因此必须选择并构建适当的docker镜像版本。
为此,在构建时通过arch_version
构建参数传递gpu
或cpu
docker build --build-arg arch_version=${ARCH_VERSION}
调用run_example.sh
时,可以通过第一个命令行参数选择适当的版本,例如
run_example.sh gpu ...
如果未指定,默认为cpu
。
配置模型和训练
要训练自己的模型,请在config.json
文件中配置实验,然后运行
python allrank/main.py --config_file_name allrank/config.json --run_id <你的实验名称> --job_dir <保存结果的位置>
训练过程的所有超参数:即模型定义、数据位置、使用的损失和指标、训练超参数等,都由config.json
文件控制。我们提供了一个模板文件config_template.json
,其中解释了支持的属性、其含义和可能的值。
请注意,按照MSLR-WEB30K约定,包含训练数据的libsvm文件应命名为train.txt
。您可以在配置中指定验证数据集的名称
(例如valid或test)。结果将保存在<job_dir>/results/<run_id>
路径下。
allRank支持使用Google Cloud Storage作为数据和作业结果的存储位置。
实现自定义损失函数
要试验自己的自定义损失,您需要实现一个函数,该函数接受两个张量(模型预测和真实标签)作为输入,并将其放在losses
包中,确保在包级别公开。
要在训练中使用它,只需在配置文件的正确位置传递函数的名称(如果您的损失方法有一些超参数,还需传递参数):
"loss": {
"name": "yourLoss",
"args": {
"arg1": val1,
"arg2: val2
}
}
应用点击模型
要应用点击模型,您首先需要训练好一个allRank模型。 然后,运行:
python allrank/rank_and_click.py --input-model-path <模型权重文件路径> --roles <要处理的数据集角色的逗号分隔列表,例如train,valid> --config_file_name allrank/config.json --run_id <你的实验名称> --job_dir <保存结果的位置>
该模型将用于对配置中指定的数据集中的所有序列进行排序。接下来,将应用配置中设置的点击模型,并将生成的点击率数据集以libSVM格式写入<job_dir>/results/<run_id>
。
然后,结果目录的路径可用作另一个allRank模型训练的输入。
持续集成
您应该运行scripts/ci.sh
来验证代码是否符合样式指南并通过单元测试。
研究
该框架是为支持研究项目Context-Aware Learning to Rank with Self-Attention而开发的。如果您在研究中使用allRank,请引用:
@article{Pobrotyn2020ContextAwareLT,
title={Context-Aware Learning to Rank with Self-Attention},
author={Przemyslaw Pobrotyn and Tomasz Bartczak and Mikolaj Synowiec and Radoslaw Bialobrzeski and Jaroslaw Bojar},
journal={ArXiv},
year={2020},
volume={abs/2005.10084}
}
此外,如果您使用NeuralNDCG损失函数,请引用相应的工作,NeuralNDCG: Direct Optimisation of a Ranking Metric via Differentiable Relaxation of Sorting:
@article{Pobrotyn2021NeuralNDCG,
title={NeuralNDCG: Direct Optimisation of a Ranking Metric via Differentiable Relaxation of Sorting},
author={Przemyslaw Pobrotyn and Radoslaw Bialobrzeski},
journal={ArXiv},
year={2021},
volume={abs/2102.07831}
}
许可证
Apache 2 许可证