设置 TorchRec

在本节中，我们将：

• 了解使用 TorchRec 的要求

• 设置环境以集成 TorchRec

• 运行基本的 TorchRec 代码

系统要求

TorchRec 仅在 AWS Linux 上进行例行测试，应该可以在类似的环境中工作。 下面演示了当前测试的兼容性矩阵：

Python 版本	3.9, 3.10, 3.11, 3.12
计算平台	中央处理器、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，请分别交换为 、 或 。cu121cpucu118cu124

通过 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 数据类型运行一个简单的前向传递：和 ：KeyedJaggedTensorEmbeddingBagCollection

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

您应该会看到带有结果嵌入向量的输出。 恭喜！您已经正确安装并运行了您的第一个 TorchRec 程序！KeyedTensor