TensorFlow 推荐系统附加组件

---

TensorFlow 推荐系统附加组件（TFRA）是一系列与大规模推荐系统相关的项目，基于 TensorFlow 构建，并引入了 动态嵌入技术，使 TensorFlow 更加适合训练 搜索、推荐和广告 模型，同时简化了复杂推荐模型的构建、评估和部署流程。详情请参阅已批准的 TensorFlow RFC #313。这些贡献将作为 TensorFlow Core 和 TensorFlow 推荐系统的补充。

对于 Apple Silicon（M1 芯片），请参考 Apple Silicon 支持。

主要特性

• 使键值数据结构（动态嵌入）可在 TensorFlow 中进行训练
• 相比静态嵌入机制，在无哈希冲突的情况下获得更好的推荐效果
• 兼容所有原生 TensorFlow 优化器和初始化器
• 兼容原生 TensorFlow CheckPoint 和 SavedModel 格式
• 完全支持在 GPU 上训练和推理推荐模型
• 支持 TF Serving 和 Triton 推理服务器 作为推理框架
• 支持多种键值实现作为动态嵌入存储，并易于扩展
  • cuckoohash_map（来自卡内基梅隆大学 Efficient Computing 团队，适用于 CPU）
  • HierarchicalKV（来自 NVIDIA，适用于 GPU）
  • Redis
• 支持基于 Horovod 的半同步训练
  • 密集权重采用同步训练
  • 稀疏权重采用异步训练

子包

• tfra.dynamic_embedding, RFC
• tfra.embedding_variable, RFC

贡献者

TensorFlow 推荐系统附加组件依赖于公众的贡献、错误修复和文档编写。本项目得以存在，离不开所有贡献者及其所在组织的支持。[贡献]

特别感谢 NVIDIA Merlin 团队以及 NVIDIA 中国开发技术支持团队，他们提供了 GPU 加速技术支持和代码贡献。

教程与演示

请参阅 教程 和 演示，以获取各子包的端到端示例。

安装

稳定版本

TensorFlow 推荐系统附加组件已在 PyPI 上提供，适用于 Linux 和 macOS。要安装最新版本，请运行以下命令：

pip install tensorflow-recommenders-addons

在 0.8 版本之前，若需安装 GPU 版本，请运行以下命令：

pip install tensorflow-recommenders-addons-gpu

使用 TensorFlow 推荐系统附加组件时：

import tensorflow as tf
import tensorflow_recommenders_addons as tfra

与 TensorFlow 的兼容性

TensorFlow C++ API 尚未稳定，因此我们仅能保证与构建 TensorFlow Recommenders-Addons (TFRA) 时所针对的版本兼容。虽然 TFRA 有可能在多个 TensorFlow 版本上运行，但也存在出现段错误或其他问题导致崩溃的风险。如果您的 TensorFlow 版本与构建时使用的版本不匹配，系统将发出警告。

此外，TFRA 自定义算子的注册没有稳定的 ABI 接口，因此即使用户的 TensorFlow 版本与我们构建时使用的版本一致，仍然需要安装兼容的 TensorFlow 环境。简单来说，TensorFlow Recommenders-Addons 的自定义算子可以与通过 pip 安装的 TensorFlow 配合使用，但如果 TensorFlow 是以不同方式编译的，则可能会出现问题。一个典型的例子就是通过 conda 安装的 TensorFlow。RFC #133 旨在解决这一问题。

兼容性矩阵

GPU 支持从版本 0.2.0 及以上开始。

TFRA	TensorFlow	编译器	CUDA	CUDNN	计算能力	CPU
0.8.0	2.16.2	GCC 8.2.1	12.3	8.9	7.0, 7.5, 8.0, 8.6, 8.9, 9.0	x86
0.8.0	2.16.2	Xcode 13.1	-	-	-	Apple M1
0.7.0	2.15.1	GCC 8.2.1	12.2	8.9	7.0, 7.5, 8.0, 8.6, 8.9, 9.0	x86
0.7.0	2.15.1	Xcode 13.1	-	-	-	Apple M1
0.6.0	2.8.3	GCC 7.3.1	11.2	8.1	6.0, 6.1, 7.0, 7.5, 8.0, 8.6	x86
0.6.0	2.6.0	Xcode 13.1	-	-	-	Apple M1
0.5.1	2.8.3	GCC 7.3.1	11.2	8.1	6.0, 6.1, 7.0, 7.5, 8.0, 8.6	x86
0.5.1	2.6.0	Xcode 13.1	-	-	-	Apple M1
0.5.0	2.8.3	GCC 7.3.1	11.2	8.1	6.0, 6.1, 7.0, 7.5, 8.0, 8.6	x86
0.5.0	2.6.0	Xcode 13.1	-	-	-	Apple M1
0.4.0	2.5.1	GCC 7.3.1	11.2	8.1	6.0, 6.1, 7.0, 7.5, 8.0, 8.6	x86
0.4.0	2.5.0	Xcode 13.1	-	-	-	Apple M1
0.3.1	2.5.1	GCC 7.3.1	11.2	8.1	6.0, 6.1, 7.0, 7.5, 8.0, 8.6	x86
0.2.0	2.4.1	GCC 7.3.1	11.0	8.0	6.0, 6.1, 7.0, 7.5, 8.0	x86
0.2.0	1.15.2	GCC 7.3.1	10.0	7.6	6.0, 6.1, 7.0, 7.5	x86
0.1.0	2.4.1	GCC 7.3.1	-	-	-	x86

更多详细信息请参阅 NVIDIA 支持矩阵。

注意

• 发布包与 TensorFlow 之间存在严格的版本绑定关系。
• 由于 TensorFlow API 的重大变化，我们只能确保版本 0.2.0 在 CPU 和 GPU 上与 TF1.15.2 兼容，但没有官方发布版本，您只能通过以下方式自行编译获得：

PY_VERSION="3.9" \
TF_VERSION="2.15.1" \
TF_NEED_CUDA=1 \
sh .github/workflows/make_wheel_Linux_x86.sh

# .whl 文件将被创建在 ./wheelhouse/

• 如果您需要使用 TensorFlow 1.14.x 或更早版本，我们建议您放弃，不过这篇文档或许能帮到您：从 TensorFlow 编译目录中提取头文件。同时，我们发现 TRFA 使用的一些算子性能更好，因此强烈建议您将 TensorFlow 更新至 2.x 版本。

从源码安装

对于所有开发者，我们推荐使用已启用 GPU 的开发容器：

docker pull tfra/dev_container:latest-tf2.15.1-python3.9  # 可用的 TensorFlow 和 Python 组合请参见 [这里](https://www.tensorflow.org/install/source#linux)
docker run --privileged --gpus all -it --rm -v $(pwd):$(pwd) tfra/dev_container:latest-tf2.15.1-python3.9

仅 CPU

您也可以从源码安装。这需要使用 Bazel 构建系统（版本 == 5.1.1）。请在编译机器上安装一个 TensorFlow，编译器需要根据已安装的 TensorFlow 知道其版本和头文件路径。

export TF_VERSION="2.15.1"  # "2.11.0" 经过充分测试。
pip install tensorflow==$TF_VERSION

git clone https://github.com/tensorflow/recommenders-addons.git
cd recommenders-addons

# 此脚本将项目与 TensorFlow 依赖项关联
python configure.py

bazel build --enable_runfiles build_pip_pkg
bazel-bin/build_pip_pkg artifacts

pip install artifacts/tensorflow_recommenders_addons-*.whl

GPU 支持

只需设置 TF_NEED_CUDA=1，其他环境变量为可选：

export TF_VERSION="2.15.1"  # "2.11.0" 经过充分测试。
export PY_VERSION="3.9" 
export TF_NEED_CUDA=1
export TF_CUDA_VERSION=12.2 # 使用 nvcc --version 检查版本
export TF_CUDNN_VERSION=8.9 # 打印("cuDNN 版本:", tf.sysconfig.get_build_info()["cudnn_version"])
export CUDA_TOOLKIT_PATH="/usr/local/cuda"
export CUDNN_INSTALL_PATH="/usr/lib/x86_64-linux-gnu"

python configure.py

然后构建并安装 pip 包：

bazel build --enable_runfiles build_pip_pkg
bazel-bin/build_pip_pkg artifacts
pip install artifacts/tensorflow_recommenders_addons_gpu-*.whl

要运行单元测试：

cp -f ./bazel-bin/tensorflow_recommenders_addons/dynamic_embedding/core/*.so ./tensorflow_recommenders_addons/dynamic_embedding/core/
pip install pytest
python tensorflow_recommenders_addons/tests/run_all_test.py
# 并运行 pytest，例如
pytest -s tensorflow_recommenders_addons/dynamic_embedding/python/kernel_tests/hkv_hashtable_ops_test.py

Apple Silicon 支持

要求：

• macOS 12.0.0+
• TensorFlow 2.15.1
• Bazel 5.1.1

通过 Pypi 在 Apple Silicon 上安装 TFRA

python -m pip install tensorflow-recommenders-addons --no-deps

从源码在 Apple Silicon 上构建 TFRA

# 安装 bazelisk
brew install bazelisk

# 从源码构建 wheel
TF_VERSION=2.15.1 TF_NEED_CUDA="0" sh .github/workflows/make_wheel_macOS_arm64.sh

# 安装 wheel 包
python -m pip install --no-deps ./artifacts/*.whl

已知问题：

TFRA 的 Apple Silicon 版本不支持以下功能：

• 数据类型 float16
• 基于 Horovod 的同步训练
• HierarchicalKV (HKV)
• save_to_file_system
• load_from_file_system
• warm_start_util

save_to_file_system 和 load_from_file_system 不受支持，因为 TFIO 在 Apple Silicon 设备上不受支持。Horovod 和 warm_start_util 也不受支持，因为原生支持的 tensorflow-macos 不支持 V1 Tensorflow 网络。

这些问题可能会在未来的版本中得到修复。

tfra.dynamic_embedding.Variable 的数据类型矩阵

值 \ 键	int64	int32	string
float	CPU, GPU	CPU, GPU	CPU
bfloat16	CPU, GPU	CPU	CPU
half	CPU, GPU	-	CPU
int32	CPU, GPU	CPU	CPU
int8	CPU, GPU	-	CPU
int64	CPU	-	CPU
double	CPU, CPU	CPU	CPU
bool	-	-	CPU
string	CPU	-	-

使用 tfra.dynamic_embedding.Variable 进行 GPU 计算

tfra.dynamic_embedding.Variable 会忽略 TensorFlow 的设备放置机制，因此您需要显式地为它指定 GPU 设备。

import tensorflow as tf
import tensorflow_recommenders_addons as tfra

de = tfra.dynamic_embedding.get_variable("VariableOnGpu",
                                         devices=["/job:ps/task:0/GPU:0", ],
                                         # ...
                                         )

GPU 使用限制

• 仅适用于具有 CUDA 计算能力 6.0 或更高版本的 NVIDIA GPU。
• 考虑到 .whl 文件的大小，目前 dim 仅支持小于或等于 200。如果需要更大的 dim，请提交问题。
• 只有 dynamic_embedding API 及其相关操作支持在 GPU 上运行。
• 对于 GPU HashTable，它们会独立管理 GPU 内存，因此需要配置 TensorFlow 以允许 GPU 内存增长，如下所示：

sess_config.gpu_options.allow_growth = True

推理

使用 TensorFlow Serving

兼容性矩阵

TFRA	TensorFlow	Serving 分支	编译器	CUDA	CUDNN	计算能力
0.8.0	2.16.2	r2.16	GCC 8.2.1	12.3	8.9	7.0、7.5、8.0、8.6、8.9、9.0
0.7.0	2.15.1	r2.15	GCC 8.2.1	12.2	8.9	7.0、7.5、8.0、8.6、8.9、9.0
0.6.0	2.8.3	r2.8	GCC 7.3.1	11.2	8.1	6.0、6.1、7.0、7.5、8.0、8.6
0.5.1	2.8.3	r2.8	GCC 7.3.1	11.2	8.1	6.0、6.1、7.0、7.5、8.0、8.6
0.5.0	2.8.3	r2.8	GCC 7.3.1	11.2	8.1	6.0、6.1、7.0、7.5、8.0、8.6
0.4.0	2.5.1	r2.5	GCC 7.3.1	11.2	8.1	6.0、6.1、7.0、7.5、8.0、8.6
0.3.1	2.5.1	r2.5	GCC 7.3.1	11.2	8.1	6.0、6.1、7.0、7.5、8.0、8.6
0.2.0	2.4.1	r2.4	GCC 7.3.1	11.0	8.0	6.0、6.1、7.0、7.5、8.0
0.2.0	1.15.2	r1.15	GCC 7.3.1	10.0	7.6	6.0、6.1、7.0、7.5
0.1.0	2.4.1	r2.4	GCC 7.3.1	-	-	-

通过 TensorFlow Serving 中的自定义操作启用 TFRA 模型。

## 如果启用 GPU 操作
export SERVING_WITH_GPU=1 

## 指定 TFRA 的分支
export TFRA_BRANCH="master" # 目前可使用 `master` 和 `r0.6` 分支。

## 创建工作区，您可以根据需要修改目录。
export TFRA_SERVING_WORKSPACE=~/tfra_serving_workspace/
mkdir -p $TFRA_SERVING_WORKSPACE && cd $TFRA_SERVING_WORKSPACE

## 根据“兼容性矩阵”克隆 Serving 和 TFRA 的发布分支。
git clone -b r2.8 https://github.com/tensorflow/serving.git
git clone -b $TFRA_BRANCH https://github.com/tensorflow/recommenders-addons.git

## 运行配置脚本
cd $TFRA_SERVING_WORKSPACE/recommenders-addons/tools
bash config_tfserving.sh $TFRA_BRANCH $TFRA_SERVING_WORKSPACE/serving $SERVING_WITH_GPU

## 构建带有 TFRA 操作的 Serving。
cd $TFRA_SERVING_WORKSPACE/serving
./tools/run_in_docker.sh bazel build tensorflow_serving/model_servers:tensorflow_model_server

更多详细信息，请参阅脚本 ./tools/config_tfserving.sh。

注意

• 分布式推理仅在使用 Redis 作为键值存储时才支持。
• 参考文档：https://www.tensorflow.org/tfx/serving/custom_op

使用 Triton

在构建自定义操作共享库时，务必使用与 Triton 中相同的 TensorFlow 版本。您可以在 Triton 发行说明 中找到 TensorFlow 版本。确保使用正确 TensorFlow 版本的简单方法是使用与 Triton 容器相对应的 NGC TensorFlow 容器。例如，如果您使用的是 Triton 23.05 版本，则应使用 TensorFlow 23.05 版本。

docker pull nvcr.io/nvidia/tritonserver:22.05-py3

export TFRA_BRANCH="master"
git clone -b $TFRA_BRANCH https://github.com/tensorflow/recommenders-addons.git
cd recommenders-addons

python configure.py
bazel build //tensorflow_recommenders_addons/dynamic_embedding/core:_cuckoo_hashtable_ops.so ##bazel 5.1.1 已经过充分测试
mkdir /tmp/so
#您也可以使用从 pip 安装包中获取的 so 文件，路径为 "(PYTHONPATH)/site-packages/tensorflow_recommenders_addons/dynamic_embedding/core/_cuckoo_hashtable_ops.so"
cp bazel-bin/tensorflow_recommenders_addons/dynamic_embedding/core/_cuckoo_hashtable_ops.so /tmp/so

#tfra saved_model 目录 "/models/model_repository"
docker run --net=host -v /models/model_repository:/models nvcr.io/nvidia/tritonserver:22.05-py3 bash -c \
  "LD_PRELOAD=/tmp/so/_cuckoo_hashtable_ops.so:${LD_PRELOAD} tritonserver --model-repository=/models/ --backend-config=tensorflow,version=2 --strict-model-config=false"

注意

• 必须设置上述 LD_LIBRARY_PATH 和 backend-config，因为默认后端是 tf1。

社区

• SIG Recommenders 邮件列表： recommenders@tensorflow.org

致谢

我们非常感谢 tensorflow/addons 的维护者们，他们借用了大量代码来自 tensorflow/addons，帮助我们构建了工作流和文档系统。同时，我们也感谢 Google 团队成员在 CI 设置和评审方面提供的帮助！

许可证

Apache License 2.0

TensorFlow Recommenders Addons (TFRA) 快速上手指南

TensorFlow Recommenders Addons (TFRA) 是基于 TensorFlow 构建的大规模推荐系统工具集。它引入了**动态嵌入（Dynamic Embedding）**技术，解决了传统静态嵌入在海量稀疏特征场景下的哈希冲突和显存限制问题，特别适用于搜索、推荐和广告场景。

1. 环境准备

系统要求

• 操作系统: Linux (x86), macOS (Intel & Apple Silicon M1/M2)
• Python 版本: 推荐 3.9 (支持 3.7+)
• 构建工具: Bazel 5.1.1 (仅源码编译时需要)

前置依赖与版本兼容性

TFRA 与 TensorFlow 版本严格绑定。请确保安装对应的 TensorFlow 版本，否则可能出现崩溃或警告。

TFRA 版本	推荐 TensorFlow 版本	GPU 支持	备注
0.8.0	2.16.2	✅ (CUDA 12.3)	最新稳定版
0.7.0	2.15.1	✅ (CUDA 12.2)	广泛使用版本
0.6.0	2.8.3 / 2.6.0	✅	-

注意：

• GPU 用户：需安装对应版本的 tensorflow (包含 GPU 支持)。
• Apple Silicon (M1/M2)：需 macOS 12.0+，TensorFlow 2.15.1，且部分功能（如 float16, Horovod, HKV）暂不支持。
• Conda 用户：由于自定义算子 ABI 接口限制，推荐使用 pip 安装的 TensorFlow，Conda 版本可能无法正常工作。

2. 安装步骤

方式一：PyPI 安装（推荐）

适用于大多数 Linux 和 macOS 用户。

标准安装 (CPU/GPU 通用，新版):

pip install tensorflow-recommenders-addons

旧版 GPU 专用 (仅限 TFRA < 0.8): 如果你使用的是 0.8 之前的版本且需要 GPU 支持：

pip install tensorflow-recommenders-addons-gpu

Apple Silicon (M1/M2) 特殊安装: 由于依赖关系复杂，建议忽略依赖自动安装，手动管理 TensorFlow：

# 先确保安装了 tensorflow-macos
python -m pip install tensorflow-recommenders-addons --no-deps

方式二：源码编译安装

适用于需要特定版本组合或使用 Conda 环境的开发者。

1. 准备环境 (以 Docker 为例，推荐):

docker pull tfra/dev_container:latest-tf2.15.1-python3.9
docker run --privileged --gpus all -it --rm -v $(pwd):$(pwd) tfra/dev_container:latest-tf2.15.1-python3.9

2. 配置与编译 (CPU 版本):

export TF_VERSION="2.15.1"
pip install tensorflow==$TF_VERSION

git clone https://github.com/tensorflow/recommenders-addons.git
cd recommenders-addons

python configure.py

bazel build --enable_runfiles build_pip_pkg
bazel-bin/build_pip_pkg artifacts

pip install artifacts/tensorflow_recommenders_addons-*.whl

3. 配置与编译 (GPU 版本): 设置 CUDA 环境变量后执行编译：

export TF_VERSION="2.15.1"
export TF_NEED_CUDA=1
export TF_CUDA_VERSION=12.2
export TF_CUDNN_VERSION=8.9
# 根据实际路径调整
export CUDA_TOOLKIT_PATH="/usr/local/cuda" 
export CUDNN_INSTALL_PATH="/usr/lib/x86_64-linux-gnu"

python configure.py
bazel build --enable_runfiles build_pip_pkg
bazel-bin/build_pip_pkg artifacts

pip install artifacts/tensorflow_recommenders_addons_gpu-*.whl

3. 基本使用

导入库

import tensorflow as tf
import tensorflow_recommenders_addons as tfra

核心功能：动态嵌入 (Dynamic Embedding)

与传统 tf.Embedding 不同，tfra.dynamic_embedding.Variable 允许嵌入表大小动态增长，无需预先设定词汇表大小，且支持海量稀疏特征。

简单示例：

import tensorflow as tf
import tensorflow_recommenders_addons as tfra

# 1. 创建动态嵌入变量
# keys 可以是 int64, string 等，values 支持多种数据类型
embedding = tfra.dynamic_embedding.Variable(
    name="user_embedding",
    key_dtype=tf.int64,
    value_dtype=tf.float32,
    initializer=tf.random_normal_initializer(),
    dim=128  # 嵌入维度
)

# 2. 定义查找操作 (Lookup)
# ids: [batch_size]
ids = tf.constant([1001, 2002, 999999], dtype=tf.int64) 

# 获取嵌入向量: [batch_size, dim]
# 如果 key 不存在，会自动创建并初始化
embeddings = embedding.lookup(ids)

# 3. 在模型中使用 (配合 GradientTape)
with tf.GradientTape() as tape:
    loss = tf.reduce_mean(embeddings ** 2) # 示例损失函数

# 4. 计算梯度并优化
grads = tape.gradient(loss, embedding.trainable_variables)
optimizer = tf.keras.optimizers.SGD(learning_rate=0.01)
optimizer.apply_gradients(zip(grads, embedding.trainable_variables))

print("动态嵌入形状:", embeddings.shape)

关键特性提示

• 设备放置：tfra.dynamic_embedding.Variable 会忽略 TensorFlow 默认的设备放置机制。若在 GPU 上训练，必须显式指定设备或在创建变量时配置存储后端（如 HierarchicalKV）。
• 保存与加载：完全兼容 TensorFlow 原生的 CheckPoint 和 SavedModel 格式。
• 分布式训练：支持基于 Horovod 的半同步训练（稠密权重同步，稀疏权重异步）。