allRank - 深入的PyTorch排序学习框架，支持多种神经网络模型

项目介绍：allRank

allRank 是一个基于 PyTorch 的框架，用于训练神经学习排序（Learning-to-Rank，简称 LTR）模型。这个框架提供了一系列常见的点级、对级和序列级损失函数，以及全连接和 Transformer 类似的评分函数。此外，它还支持常用的评估指标，如归一化折现累积分（NDCG）和平均互惠排序（MRR），并提供用于模拟点击率数据的点击模型。

项目的动机

allRank 致力于为用户提供一个简单且灵活的平台，以便在各种 LTR 神经网络模型和损失函数上进行实验。框架便于添加自定义的损失函数，并能灵活配置模型和训练过程。我们希望 allRank 能同时推动神经 LTR 的学术研究和工业应用的发展。

功能特性

已实现的损失函数

ListNet（适用于二进制和分级相关性）
ListMLE
RankNet
序列损失
LambdaRank
LambdaLoss
ApproxNDCG
均方根误差（RMSE）
NeuralNDCG

使用指南

为了帮助用户快速开始，我们提供了 run_example.sh 脚本。这个脚本会生成一个模拟的排名数据集，并基于该数据集训练一个 Transformer 模型。模拟数据将存储在 dummy_data 目录中，实验结果将保存在 test_run 目录内。需要 Docker 才能运行这些示例。

选择适合的架构版本（GPU 与 CPU）

由于 torch 的二进制文件在 GPU 和 CPU 中有所不同（GPU 版本无法在 CPU 上运行），用户需要选择并构建合适的 Docker 镜像版本。在构建时，需要通过 build-arg 传递 gpu 或 cpu 参数。

运行示例时，可以通过第一个命令行参数指定所需的版本，例如：

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）可以通过 config 制定。结果会保存到路径 <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 <结果保存路径>

该模型将用于排名配置中指定的数据集中的所有组。接下来，配置中设定的点击模型将被应用，结果的点击率数据集将被写入 <job_dir>/results/<run_id> 中的 libSVM 格式。随后，结果目录的路径可用于训练另一个 allRank 模型。

持续集成

用户应运行 scripts/ci.sh 来验证代码是否通过了风格指导和单元测试。

学术研究

allRank 框架用于支持“基于上下文的自注意力学习排序”研究项目。如果在研究中使用 allRank，请引用相关的论文。此外，如果使用了 NeuralNDCG 损失函数，请引用对应的研究工作。

许可证

allRank 采用 Apache 2 许可证。