datawhalechina/torch-rechub

What is it?

Key features

• 01Implements 30+ ranking, matching, multi-task and generative recommendation models.
• 02Provides unified trainers and data loaders for CTR, matching and multi-task training.
• 03Exports trained models to ONNX for production deployment.
• 04Runs on CPU, NVIDIA GPU, AMD ROCm and Huawei Ascend NPU.
• 05Integrates WandB, SwanLab and TensorBoardX for experiment tracking.
• 06Offers optional FAISS, Annoy and Milvus support for vector retrieval.

Use cases

• Prototyping CTR prediction models on datasets such as Criteo or Avazu.
• Building two-tower retrieval models and exporting user and item towers separately.
• Training multi-task ranking models for joint objectives like click and conversion.
• Experimenting with LLM-based or generative recommendation methods such as HSTU and HLLM.

Quick start

1. 1

   Install PyTorch and the package with pip

   pip install torch && pip install torch-rechub
2. 2

   Clone the repository

   git clone https://github.com/datawhalechina/torch-rechub.git
3. 3

   Install dependencies with uv

   uv sync
4. 4

   Run the DSSM matching example from its folder

   cd examples/matching && python run_ml_dssm.py

Things to know

• !Requires a PyTorch build matched to your hardware, driver and runtime, which can take setup effort.
• !Some features such as ONNX export, visualization and tracking need optional extras installed separately.

Topics

FAQ

What is torch-rechub?

A PyTorch framework for building recommendation models, with 30+ built-in ranking, matching, multi-task and generative models and ONNX export. Torch-RecHub is a PyTorch-based framework for building recommender systems. It bundles a library of classic and recent recommendation models, including CTR ranking models such as DeepFM and DCN, two-tower and sequential matching models such as DSSM and SASRec, multi-task models such as MMoE and PLE, and generative recommenders such as HSTU and TIGER. Each model is trained through shared trainers for CTR, matching and multi-task setups, with unified data loading and evaluation workflows.

Is torch-rechub free to use? What license does it have?

torch-rechub is open source under the MIT license. Check the license terms before using it in commercial products.

Is torch-rechub still maintained?

Actively maintained. The last commit to torch-rechub was on Oct 22, 2026.

What can I use torch-rechub for?

Prototyping CTR prediction models on datasets such as Criteo or Avazu. Building two-tower retrieval models and exporting user and item towers separately. Training multi-task ranking models for joint objectives like click and conversion. Experimenting with LLM-based or generative recommendation methods such as HSTU and HLLM.

How do I install torch-rechub?

Install PyTorch and the package with pip: pip install torch && pip install torch-rechub Clone the repository: git clone https://github.com/datawhalechina/torch-rechub.git Install dependencies with uv: uv sync Run the DSSM matching example from its folder: cd examples/matching && python run_ml_dssm.py

Where can I get torch-rechub?

The source code and installation instructions are on GitHub: https://github.com/datawhalechina/torch-rechub