探索大规模推荐系统的新里程：TorchRec（Beta 版）

最新推荐文章于 2024-10-08 09:30:05 发布

屈游会

最新推荐文章于 2024-10-08 09:30:05 发布

阅读量413

点赞数 4

本文链接：https://blog.csdn.net/gitblog_00417/article/details/141121621

版权

探索大规模推荐系统的新里程：TorchRec（Beta 版）

torchrecPytorch domain library for recommendation systems项目地址:https://gitcode.com/gh_mirrors/to/torchrec

Tech Docs

在现代的互联网世界中，推荐系统已经成为用户体验的核心组成部分。TorchRec，一个基于PyTorch的领域库，专为构建大型推荐系统（RecSys）而设计，它提供了必要的稀疏性和并行性原语，以处理大规模的嵌入表。TorchRec允许开发者在多个GPU上训练模型，实现数据并行和模型并行的混合优化。

项目核心功能

TorchRec包括以下组件：

并行主义原语：让编写大型、高性能的多设备/多节点模型变得更加简单。
TorchRec分片器：支持数据并行、表格智能、行智能等多样化分片策略，灵活处理嵌入表。
TorchRec规划器：自动为模型生成优化的分片计划。
流水线训练：通过重叠数据加载、设备传输、计算操作，提高训练性能。
优化内核：针对RecSys进行优化，由FBGEMM提供动力。
量化支持：实现降低精度的训练和推理。
RecSys通用模块：标准模块化设计，方便复用。
生产验证的模型架构：经过实际场景考验的高效模型设计。
RecSys数据集：如Criteo点击日志和MovieLens，可用于测试和开发。
端到端训练示例：例如使用Criteo点击日志数据集训练的dlrm事件预测模型。

安装

TorchRec要求Python版本大于等于3.8，以及CUDA版本大于等于11.8（虽然高度推荐但非必需）。以下是使用Python 3.8和CUDA 12.1的安装示例。对于CPU或CUDA 11.8，请相应地将"cu121"替换为"cpu"或"cu118"。

安装步骤

# Nightly版本
pip install torch --index-url https://download.pytorch.org/whl/nightly/cu121
pip install fbgemm-gpu --index-url https://download.pytorch.org/whl/nightly/cu121
pip install torchmetrics==1.0.3
pip install torchrec --index-url https://download.pytorch.org/whl/nightly/cu121

# Stable版本 (via pytorch.org)
pip install torch --index-url https://download.pytorch.org/whl/cu121
pip install fbgemm-gpu --index-url https://download.pytorch.org/whl/cu121
pip install torchmetrics==1.0.3
pip install torchrec --index-url https://download.pytorch.org/whl/cu121

# Stable版本 (仅限CUDA 12.1, 通过PyPI)
pip install torch
pip install fbgemm-gpu
pip install torchrec

也可以直接从GitHub源代码编译安装，详细步骤见项目文档。