设置 TorchRec

在本节中，我们将：

• 理解使用TorchRec所需的要求

• 设置一个环境来集成TorchRec

• 运行基本的 TorchRec 代码

系统要求

TorchRec 仅在 AWS Linux 上进行常规测试，并且应该能在类似的环境中工作。以下展示了当前正在测试的兼容性矩阵：

Python 版本	3.9, 3.10, 3.11, 3.12
计算平台	CPU, CUDA 11.8, CUDA 12.1, CUDA 12.4

除了这些要求，TorchRec的核心依赖是PyTorch和FBGEMM。如果你的系统兼容这两个库，那么TorchRec应该足够使用。

• PyTorch需求

• FBGEMM的要求

兼容性版本

TorchRec 和 FBGEMM 有匹配的版本号，在发布时一起进行测试。

• TorchRec 1.0 兼容 FBGEMM 1.0

• TorchRec 0.8 兼容 FBGEMM 0.8

• TorchRec 0.8 可能不兼容 FBGEMM 0.7

此外，TorchRec 和 FBGEMM 只会在 PyTorch 发布新版本时才发布。 因此，TorchRec 和 FBGEMM 的特定版本应对应一个特定的 PyTorch 版本：

• TorchRec 1.0 兼容 PyTorch 2.5

• TorchRec 0.8 兼容 PyTorch 2.4

• TorchRec 0.8 可能不兼容 PyTorch 2.3

安装

以下是CUDA 12.1的安装示例。对于CPU，CUDA 11.8或CUDA 12.4，请分别将cu121替换为cpu, cu118, 或cu124。

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

稳定通过PyPI（仅限CUDA 12.4）

pip install torch
pip install fbgemm-gpu
pip install torchrec

夜间版

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

从源代码构建

您还可以通过源代码构建TorchRec，以使用TorchRec中的最新更改进行开发。要从源代码构建，请参阅此参考。

运行一个简单的TorchRec示例

现在我们已经正确设置了TorchRec，让我们运行一些TorchRec代码！ 下面，我们将使用TorchRec数据类型：KeyedJaggedTensor和EmbeddingBagCollection进行一个简单的前向传播：

import torch

import torchrec
from torchrec.sparse.jagged_tensor import JaggedTensor, KeyedJaggedTensor

ebc = torchrec.EmbeddingBagCollection(
    device="cpu",
    tables=[
        torchrec.EmbeddingBagConfig(
            name="product_table",
            embedding_dim=16,
            num_embeddings=4096,
            feature_names=["product"],
            pooling=torchrec.PoolingType.SUM,
        ),
        torchrec.EmbeddingBagConfig(
            name="user_table",
            embedding_dim=16,
            num_embeddings=4096,
            feature_names=["user"],
            pooling=torchrec.PoolingType.SUM,
        )
    ]
)

product_jt = JaggedTensor(
    values=torch.tensor([1, 2, 1, 5]), lengths=torch.tensor([3, 1])
)
user_jt = JaggedTensor(values=torch.tensor([2, 3, 4, 1]), lengths=torch.tensor([2, 2]))

# Q1: How many batches are there, and which values are in the first batch for product_jt and user_jt?
kjt = KeyedJaggedTensor.from_jt_dict({"product": product_jt, "user": user_jt})

print("Call EmbeddingBagCollection Forward: ", ebc(kjt))

将上述代码保存到一个名为 torchrec_example.py 的文件中。然后，你应该能够从终端执行它：

python torchrec_example.py

你应该看到输出KeyedTensor，结果的嵌入。 恭喜！你正确地安装并运行了你的第一个TorchRec程序！