datawhalechina/torch-rechub
A Lighting Pytorch Framework for Recommendation Models, Easy-to-use and Easy-to-extend.
What it solves
Torch-RecHub 是一個輕量級的 PyTorch 框架，旨在簡化生產級推薦系統的建立。它消除從頭實作複雜推薦演算法的工程負擔，讓開發者只需幾行程式碼即可建構與部署模型。
How it works
此框架提供模組化的推薦任務管線，包含標準化的資料載入、訓練與評估。它利用 PyTorch 在多種後端（CPU、NVIDIA CUDA、AMD ROCm、華為 Ascend NPU）上進行硬體加速。系統依任務類型劃分為專門的 trainer：CTRTrainer 用於排序/CTR 預測，MatchTrainer 用於匹配，MTLTrainer 用於多任務學習。
Who it’s for
適用於需要快速實作與測試推薦模型的機器學習工程師與研究人員，也適合希望透過 ONNX 匯出將模型部署到生產環境的開發者。
Highlights
- Extensive Model Library: 包含超過 30 種主流模型，涵蓋排序（如 DeepFM、DIN）、匹配（如 DSSM、SASRec）、多任務（如 MMoE、PLE）以及生成式推薦（如 HSTU、TIGER）。
- Production-Ready Deployment: 支援一鍵匯出為 ONNX 格式，便於在生產環境中無縫整合。
- Broad Hardware Support: 相容多種 GPU 架構，包括 NVIDIA、AMD 與華為 Ascend。
- Big Data Integration: 支援基於 PySpark 的資料處理，並可與向量資料庫如 FAISS 與 Milvus 整合以提供檢索服務。
- Experiment Tracking: 內建與 WandB、SwanLab、TensorBoardX 的整合，用於視覺化與監控。