在线文档: https://datawhalechina.github.io/torch-rechub/

Torch-RecHub —— 用10行代码构建生产级推荐系统。30+主流模型开箱即用，一键导出 ONNX 格式，让您专注于业务而非工程实现。

✨ 特性

• 模块化设计: 轻松添加新模型、数据集和评估指标。
• 基于 PyTorch: 充分利用 PyTorch 的动态图和 GPU 加速能力。支持 NVIDIA GPU 和华为 Ascend NPU。
• 丰富的模型库: 涵盖 30+ 经典及前沿推荐算法（匹配、排序、多任务、生成式推荐等）。
• 标准化流程: 提供统一的数据加载、训练和评估工作流。
• 简单配置: 通过配置文件或命令行参数调整实验设置。
• 可复现性: 设计确保实验结果可复现。
• ONNX 导出: 将训练好的模型导出为 ONNX 格式，便于无缝部署到生产环境。
• 跨引擎数据处理: 支持基于 PySpark 的数据处理与转换，方便在大数据管道中部署。
• 实验可视化与跟踪: 内置对 WandB、SwanLab 和 TensorBoardX 的统一集成。

📖 目录

• 🔥 Torch-RecHub - 一款轻量、高效、易用的 PyTorch 推荐框架
  • ✨ 特性
  • 📖 目录
  • 🔧 安装
    • 要求
    • 安装步骤
  • 🚀 快速入门
  • 📂 项目结构
  • 💡 支持的模型
  • 📊 支持的数据集
  • 🧪 示例
    • 排序（CTR 预测）
    • 多任务排序
    • 匹配模型
    • 模型可视化
  • 👨‍💻‍ 贡献者
  • 🤝 贡献
  • 📜 许可证
  • 📚 引用
  • 📫 联系方式
  • ⭐️ 星标历史

🔧 安装

要求

• Python 3.9+
• PyTorch 1.7+（建议使用支持 CUDA 的版本以获得 GPU 加速）
• NumPy
• Pandas
• SciPy
• Scikit-learn

安装步骤

稳定版（推荐）:

# 根据设备安装对应的 PyTorch
pip install torch                                                     # CPU
pip install torch --index-url https://download.pytorch.org/whl/cu121  # GPU (CUDA 12.1)
pip install torch torch-npu                                           # NPU (华为 Ascend，需 torch-npu >= 2.5.1)

pip install torch-rechub

最新版:

# 如果尚未安装 uv，先安装
pip install uv

# 克隆并安装
git clone https://github.com/datawhalechina/torch-rechub.git
cd torch-rechub

# 根据设备安装对应的 PyTorch
uv pip install torch                                                     # CPU
uv pip install torch --index-url https://download.pytorch.org/whl/cu121  # GPU (CUDA 12.1)
uv pip install torch torch-npu                                           # NPU (华为 Ascend，需 torch-npu >= 2.5.1)

uv sync

可选依赖

可通过 uv sync --extra <name> 或 pip install "torch-rechub[<name>]" 安装额外的依赖组。

• annoy: 添加基于 Annoy 的近似最近邻索引，用于检索服务。
• faiss: 添加基于 FAISS 的向量索引，适用于高性能检索实验。
• milvus: 添加 Milvus 客户端支持，便于对接外部向量数据库服务流程。
• bigdata: 添加 PyArrow 支持，用于 Parquet 格式的数据加载和大数据预处理。
• onnx: 添加 ONNX 导出、推理运行时及模型转换相关依赖。
• visualization: 添加使用 TorchView 和 Graphviz 的模型图可视化支持。
• tracking: 添加对 WandB、SwanLab 和 TensorBoardX 的集成，用于实验跟踪。
• dev: 添加测试、代码风格检查、类型注解及本地开发工具。

🚀 快速入门

以下是一个在 MovieLens 数据集上训练 DSSM 模型的简单示例：

# 如果使用最新版，先克隆仓库
git clone https://github.com/datawhalechina/torch-rechub.git
cd torch-rechub
uv sync

# 运行匹配示例（需先进入脚本目录，因为脚本使用相对路径加载数据）
cd examples/matching
python run_ml_dssm.py

# 或者使用自定义参数：
python run_ml_dssm.py --model_name dssm --device 'cuda:0' --learning_rate 0.001 --epoch 50 --batch_size 4096 --weight_decay 0.0001 --save_dir 'saved/dssm_ml-100k'

# 运行排序示例
cd ../ranking
python run_criteo.py

训练完成后，模型文件将保存在 saved/dssm_ml-100k 目录下（或您配置的目录）。

📂 项目结构

torch-rechub/             # 根目录
├── README.md             # 项目文档
├── pyproject.toml        # 项目配置与依赖
├── torch_rechub/         # 核心库
│   ├── basic/            # 基础组件
│   │   ├── activation.py # 激活函数
│   │   ├── features.py   # 特征工程
│   │   ├── layers.py     # 神经网络层
│   │   ├── loss_func.py  # 损失函数
│   │   └── metric.py     # 评估指标
│   ├── models/           # 推荐模型实现
│   │   ├── matching/     # 匹配模型（DSSM/MIND/GRU4Rec等）
│   │   ├── ranking/      # 排序模型（WideDeep/DeepFM/DIN等）
│   │   └── multi_task/   # 多任务模型（MMoE/ESMM等）
│   ├── trainers/         # 训练框架
│   │   ├── ctr_trainer.py    # CTR 预测训练器
│   │   ├── match_trainer.py  # 匹配模型训练器
│   │   └── mtl_trainer.py    # 多任务学习训练器
│   └── utils/            # 工具函数
│       ├── data.py       # 数据处理工具
│       ├── match.py      # 匹配相关工具
│       ├── mtl.py        # 多任务相关工具
│       └── onnx_export.py # ONNX 导出工具
├── examples/             # 示例脚本
│   ├── matching/         # 匹配任务示例
│   ├── ranking/          # 排序任务示例
│   └── generative/       # 生成式推荐示例（HSTU、HLLM 等）
├── docs/                 # 文档（VitePress，多语言）
├── tutorials/            # Jupyter 教程
├── tests/                # 单元测试
├── config/               # 配置文件
└── scripts/              # 工具脚本

💡 支持的模型

该框架目前支持 30+ 种主流推荐模型：

排序模型 - 13

模型	论文	描述
DeepFM	IJCAI 2017	FM + Deep 联合训练
Wide&Deep	DLRS 2016	记忆 + 泛化
DCN	KDD 2017	显式特征交叉
DCN-v2	WWW 2021	增强型交叉网络
DIN	KDD 2018	用户兴趣注意力
DIEN	AAAI 2019	兴趣演化建模
BST	DLP-KDD 2019	Transformer 用于序列
AFM	IJCAI 2017	注意力加权 FM
AutoInt	CIKM 2019	自动特征交互学习
FiBiNET	RecSys 2019	特征重要性 + 双线性
DeepFFM	RecSys 2019	字段感知 FM
EDCN	KDD 2021	增强版 DCN

匹配模型 - 12

模型	论文	描述
DSSM	CIKM 2013	经典双塔模型
YoutubeDNN	RecSys 2016	YouTube 深度检索
YoutubeSBC	RecSys 2019	采样偏差校正
MIND	CIKM 2019	多兴趣动态路由
SINE	WSDM 2021	稀疏兴趣网络
GRU4Rec	ICLR 2016	GRU 用于序列
SASRec	ICDM 2018	自注意力序列
NARM	CIKM 2017	神经注意会话
STAMP	KDD 2018	短期注意力记忆优先
ComiRec	KDD 2020	可控多兴趣

多任务模型 - 5

模型	论文	描述
ESMM	SIGIR 2018	整体空间多任务
MMoE	KDD 2018	多门混合专家
PLE	RecSys 2020	渐进式分层提取
AITM	KDD 2021	自适应信息迁移
SharedBottom	-	经典共享底部

生成式推荐 - 3

模型	论文	描述
HSTU	Meta 2024	层次化序列转换单元，驱动 Meta 的万亿参数推荐系统
HLLM	2024	层次化 LLM 用于推荐，结合 LLM 语义理解
TIGER	NeurIPS 2023	基于 T5 的生成式检索用于推荐，并生成语义 ID

📊 支持的数据集

该框架为以下常用数据集提供了内置支持或预处理脚本：

• MovieLens
• Amazon
• Criteo
• Avazu
• Census-Income
• BookCrossing
• Ali-ccp
• Yidian
• ...

预期的数据格式通常是一个包含以下内容的交互文件：

• 用户ID
• 项目ID
• 评分（可选）
• 时间戳（可选）

有关具体格式要求，请参阅 tutorials 目录中的示例代码。examples/ 目录已在每个场景子目录中包含了示例数据集，您可以直接使用这些数据集进行快速实验和调试。

您可以通过确保您的数据集符合框架的数据格式要求，或者编写自定义数据加载器，轻松集成自己的数据集。

🧪 示例

所有模型使用示例均可在 /examples 中找到。

排序（CTR 预测）

from torch_rechub.models.ranking import DeepFM
from torch_rechub.trainers import CTRTrainer
from torch_rechub.utils.data import DataGenerator

dg = DataGenerator(x, y)
train_dataloader, val_dataloader, test_dataloader = dg.generate_dataloader(split_ratio=[0.7, 0.1], batch_size=256)

model = DeepFM(deep_features=deep_features, fm_features=fm_features, mlp_params={"dims": [256, 128], "dropout": 0.2, "activation": "relu"})

ctr_trainer = CTRTrainer(model)
ctr_trainer.fit(train_dataloader, val_dataloader)
auc = ctr_trainer.evaluate(ctr_trainer.model, test_dataloader)
ctr_trainer.export_onnx("deepfm.onnx")

多任务排序

from torch_rechub.models.multi_task import SharedBottom, ESMM, MMOE, PLE, AITM
from torch_rechub.trainers import MTLTrainer

task_types = ["classification", "classification"] 
model = MMOE(features, task_types, 8, expert_params={"dims": [32,16]}, tower_params_list=[{"dims": [32, 16]}, {"dims": [32, 16]}])

mtl_trainer = MTLTrainer(model)
mtl_trainer.fit(train_dataloader, val_dataloader)
auc = ctr_trainer.evaluate(ctr_trainer.model, test_dataloader)
mtl_trainer.export_onnx("mmoe.onnx")

匹配模型

from torch_rechub.models.matching import DSSM
from torch_rechub.trainers import MatchTrainer
from torch_rechub.utils.data import MatchDataGenerator

dg = MatchDataGenerator(x, y)
train_dl, test_dl, item_dl = dg.generate_dataloader(test_user, all_item, batch_size=256)

model = DSSM(user_features, item_features, temperature=0.02,
             user_params={
                 "dims": [256, 128, 64],
                 "activation": 'prelu',  
             },
             item_params={
                 "dims": [256, 128, 64],
                 "activation": 'prelu', 
             })

match_trainer = MatchTrainer(model)
match_trainer.fit(train_dl)
match_trainer.export_onnx("dssm.onnx")
# 对于双塔模型，您可以分别导出用户塔和物品塔：
# match_trainer.export_onnx("user_tower.onnx", mode="user")
# match_trainer.export_onnx("item_tower.onnx", mode="item")

模型可视化

# 可视化模型架构（需：pip install torch-rechub[visualization]）
graph = ctr_trainer.visualization(depth=4)  # 生成计算图
ctr_trainer.visualization(save_path="model.pdf", dpi=300)  # 保存为高分辨率 PDF

👨‍💻‍ 贡献者

感谢所有贡献者！

🤝 贡献

我们欢迎任何形式的贡献！请参阅 CONTRIBUTING.md 以获取详细的贡献指南。

我们也欢迎您通过 Issues 提交错误报告和功能建议。

📜 许可证

本项目采用 MIT 许可证 许可。

📚 引用

如果您在研究或工作中使用了本框架，请考虑引用：

@misc{torch_rechub,
    title = {Torch-RecHub},
    author = {Datawhale},
    year = {2022},
    publisher = {GitHub},
    journal = {GitHub repository},
    howpublished = {\url{https://github.com/datawhalechina/torch-rechub}},
    note = {一个基于 PyTorch 的推荐系统框架，提供易于使用且可扩展的解决方案}
}

📫 联系方式

• 项目负责人: 1985312383
• GitHub Discussions

⭐️ 星标历史

---

最后更新：2026年3月20日

Torch-RecHub 快速上手指南

Torch-RecHub 是一个轻量级、高效且易用的 PyTorch 推荐系统框架。它内置了 30+ 主流推荐模型，支持一键 ONNX 部署，旨在让开发者只需少量代码即可构建生产级推荐系统。

1. 环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统: Linux, macOS 或 Windows
• Python: 3.9 及以上版本
• PyTorch: 1.7 及以上版本（推荐使用支持 CUDA 的版本以加速训练）
• 核心依赖: NumPy, Pandas, SciPy, Scikit-learn
• 硬件支持: 支持 NVIDIA GPU 及华为昇腾 NPU

2. 安装步骤

方案 A：安装稳定版（推荐）

直接使用 pip 安装。建议先根据硬件环境安装对应的 PyTorch 版本。

# 1. 安装 PyTorch (根据硬件选择其一)
pip install torch                                                     # CPU 版本
pip install torch --index-url https://download.pytorch.org/whl/cu121  # GPU 版本 (CUDA 12.1)
pip install torch torch-npu                                           # NPU 版本 (华为昇腾，需 torch-npu >= 2.5.1)

# 2. 安装 Torch-RecHub
pip install torch-rechub

国内加速提示：如果遇到下载速度慢的问题，可使用清华源或阿里源： pip install torch-rechub -i https://pypi.tuna.tsinghua.edu.cn/simple

方案 B：安装最新版（开发版）

如果您需要体验最新功能或贡献代码，可以通过源码安装。

# 1. 安装 uv 工具 (如未安装)
pip install uv

# 2. 克隆项目
git clone https://github.com/datawhalechina/torch-rechub.git
cd torch-rechub

# 3. 安装 PyTorch (根据硬件选择其一)
uv pip install torch                                                     # CPU
uv pip install torch --index-url https://download.pytorch.org/whl/cu121  # GPU
uv pip install torch torch-npu                                           # NPU

# 4. 同步依赖并安装
uv sync

可选依赖

根据具体需求安装额外功能模块（如向量检索、大数据处理、可视化等）：

# 示例：安装 FAISS 向量索引和 ONNX 导出功能
pip install "torch-rechub[faiss,onnx]"
# 或使用 uv
uv sync --extra faiss --extra onnx

支持的功能组包括：annoy, faiss, milvus, bigdata, onnx, visualization, tracking, dev。

3. 基本使用

以下以在 MovieLens 数据集上训练 DSSM（双塔模型）为例，展示最简使用流程。

步骤 1：获取示例代码

如果您安装的是稳定版，需先克隆仓库以获取示例脚本和数据路径配置：

git clone https://github.com/datawhalechina/torch-rechub.git
cd torch-rechub
# 如果之前未安装依赖，在此处执行 uv sync 或 pip install -e .

步骤 2：运行匹配模型示例

进入匹配任务示例目录并运行脚本：

cd examples/matching
python run_ml_dssm.py

步骤 3：自定义参数训练

您也可以通过命令行参数灵活调整超参数（如学习率、批次大小、设备类型等）：

python run_ml_dssm.py \
    --model_name dssm \
    --device 'cuda:0' \
    --learning_rate 0.001 \
    --epoch 50 \
    --batch_size 4096 \
    --weight_decay 0.0001 \
    --save_dir 'saved/dssm_ml-100k'

训练完成后，模型文件将自动保存至指定的 save_dir 目录中。

其他任务示例

• 排序任务 (CTR 预估):

  cd ../ranking
  python run_criteo.py
  

• 多任务学习: 参考 examples/multi_task 目录。
• 生成式推荐: 参考 examples/generative 目录。

更多详细用法、模型列表及 API 文档请访问：https://datawhalechina.github.io/torch-rechub/