tfra中的两种 embedding 稀疏参数的使用方法

5,625次阅读

共计 6674 个字符，预计需要花费 17 分钟才能阅读完成。

推荐系统是机器学习的重要应用领域，能够根据用户偏好自动推送相关内容，比如展示商品，投放广告，推荐视频、新闻等多媒体内容等。

在推荐系统领域，Embedding 已成为处理 ID 类稀疏特征的常用手段，作为一种“函数映射”，Embedding 通常将高维稀疏特征映射为低维稠密向量，再进行模型端到端训练。

在 TensorFlow 框架中，全部以稠密 Tensor 为基本数据单元对数据进行计算、存储以及传输。TensorFlow 也基于稠密 Tensor 提供了静态 Embedding 机制，用于存储 Embedding 的 Tensor shape 固定为 [vocabulary_size, embedding_dimension]，vocabulary_size 通常由 ID 空间决定，需要将其 Hash 映射到 vocabulary_size 范围内，在大规模推荐场景下，静态 Embedding 机制存在以下弊端：

特征冲突

特征 IDs 通常为字符串类型且规模庞大，这会导致初始 vocabulary_size 难以估计，如果 vocabulary_size 估计过小，则导致 Hash 冲突率增加，即不同特征可能查找到相同的 Embedding，造成特征冲突，影响模型效果：

tfra中的两种 embedding 稀疏参数的使用方法

内存浪费

如果 vocabulary_size 估计过大，则导致 Variable 会为大量永远不会被查找的 Embedding 预留资源，即内存浪费

在线学习不友好

在线学习场景中，随着训练持续进行，新的特征不断增加，老的特征要不断淘汰，这与稠密 Tensor 的形状必须事先确定相矛盾：特征不能随训练任意增加，通过算法淘汰某些不重要的特征时，也无法单独释放该特征 embedding 占用的内存资源

低效的 IO

推荐模型的更新过程具有稀疏性，即一段时间内训练更新的参数只占总量很少的一部分，但静态 Embedding 机制下，模型存取需要处理整个稠密 Tensor，这会带来极大的 IO 开销，难以支持超大模型训练

过去几年，工业界针对上述问题，在 TensorFlow 支持大规模推荐模型方面进行了大量探索，今天，我们很荣幸地推出 TensorFlow Recommenders-Addons (TFRA) 134，这款开源 TensorFlow 软件包集成了相关优秀成果，使 TensorFlow 能够以更原生、更高效的方式支持 ID 类推荐模型的训练。相关成果已经过多家互联网公司真实商业场景的检验，实践证明，这些努力和进展大幅改善了这些企业实际业务的推荐效果，带来了真实的商业收益。同时我们也通过 SIG-Recommenders 维护并回馈给业界。

TFRA 目前包含两个独立组件，其中一个是 DynamicEmbedding 组件，一个是 EmbeddingVariable 组件，其实这两个组件干的应该是同一件事情，就是解决大规模模型训练下稀疏参数的训练问题。

下面我们分别介绍这两个组件的原理和使用：

tfra.dynamic_embedding 组件基于腾讯微信看一看团队戎海栋、张亚霏、程川等人 2020 年提出的稀疏域隔离方案 109 实现，其设计目标如下：

使 TensorFlow 可以训练 Keys-Values 数据结构(Hash Table)
与 TensorFlow 原有功能有更好的兼容性，不改变算法工程师建模习惯

该组件具有如下优点：

有效避免 Hash 冲突问题，保证特征无损，能有效提升模型效果
内存动态伸缩，训练更省资源
复用 TensorFlow 所有优化器和初始化器
API 与原 TF 的 tf.nn.embedding_lookup 等 API 同名且行为一致，学习成本低

2.1.1 分层表达稀疏参数

推荐模型的稀疏参数通常存于 HashTable 中，但是 TensorFlow 优化器是为训练稠密 Tensor 而设计的，无法直接训练 HashTable 数据，幸运的是，我们观察到推荐类模型的训练过程具有稀疏性，即：每个迭代步所训练的参数仅是整个模型极少的一部分，因此我们借助 TensorFlow ResourceVariable 机制使用稠密 Tensor 表达“迭代参数”：

tfra中的两种 embedding 稀疏参数的使用方法

类似于文件系统的 memory cache 机制分两层表达稀疏参数

用原生 tf.variable/tf.tensor 表达达迭代参数 (iteration weigths)
用 tf.lookup.MutableHashTable 保存全量参数(full weights)

为了保持训练的连贯性，我们还需要做的是前向计算前从全量参数中读取本次迭代用到的 embedding 参数，反向计算后把训练好的迭代参数写回全量参数。

tfra中的两种 embedding 稀疏参数的使用方法

2.1.2 迭代参数继承自 ResourceVariable

为了保证与 TensorFlow 高度兼容性，我们需要迭代参数是个可被 TensorFlow 原生优化器训练的 ResourceVariable，同时也需要扩展一些方法以便支持稀疏训练有关的能力——为此我们定义了 TrainableWrapper 类，它是 ResourceVariable 的子类：

扩展了 prefetch 和 update 方法实现对 HashTable 读写操作
持有 HashTable(de.Variable) 和特征 ID(ids) 对象以支持前向读取 、反向更新

tfra中的两种 embedding 稀疏参数的使用方法

收益：原生优化器可以直接训练 TrainableWrapper

2.1.3 embedding_lookup 同名 API

tfra中的两种 embedding 稀疏参数的使用方法

embedding_lookup API 是迭代参数与全量参数间的桥梁

embedding lookup 是 Embedding 技术的关键语义，它通过查找操作将每个迭代步需要处理的特征权重从海量稀疏参数中提取出来，我们设计了和 TensorFlow 静态 Embedding 机制同名的 API：embedding_lookup 23，它是整个 pipeline 的信息汇聚点，我们巧妙的利用这个 API 创建了 TrainableWrapper——它具有两个关键作用：

维护稀疏参数 HashTable(params) 和 IDs（样本特征值）的关系
从 HashTable 查找中最新稀疏参数，并更新到其持有的内存资源中

这样做的好处是 API 设计更加自然，同时 API 语义和行为都和 TensorFlow 静态 Embedding 同名 API 相同，有效降低算法工程师的学习成本。

2.1.4 API 概览

相关 API 均定义在 tfra.dynamic_embedding 命名空间下（相关链接见文末）：

get_variable(…) 23：定义一个动态 Embedding 的Variable，语义与tf.get_variable相同
embedding_lookup(…) 23：动态版本的 tf.nn.embedding_lookup，参数行为相同
embedding_lookup_sparse(…) 7：动态版本的 tf.nn.embedding_lookup_sparse 7
safe_embedding_lookup_sparse(…) 8：动态版本的 tf.nn.safe_embedding_lookup_sparse 4
DynamicEmbeddingOptimizer(…) 5：原生优化器修饰器，使他们支持动态 Embedding 训练

完整定义请访问：tfra.dynamic_embedding API Docs 45

安装 TFRA 软件包

pip install tensorflow-recommenders-addons

导入 TFRA 的 dynamic_embedding 组件

import tensorflow as tf
from tensorflow_recommenders_addons import dynamic_embedding as de

创建一个动态 Embedding 变量，其中 initializer 可复用 TensorFlow 全部原生初始化器

w = de.get_variable(name="dynamic_embeddings",
                    devices=[
                        "/job:ps/replica:0/task:0/CPU:0",
                        "/job:ps/replica:0/task:1/CPU:0"
                    ],
                    initializer=tf.random_normal_initializer(0, 0.005),
                    dim=1,
                    init_size=10000000
                   )                 )

调用 embedding_lookup API 得到 embedding（此API 与 TF 同名原生 APIs 有相同入参和行为，sparse 版本的 embedding_lookup_sparse 和 safe_embedding_lookup_sparse API 说明可参考 API docs 45）：

embedding = de.embedding_lookup(params=w, ids=x, name="sparse-weights")

使用 TensorFlow 任意优化器训练:

optimizer = de.DynamicEmbeddingOptimizer(tf.compat.v1.train.AdamOptimizer(0.001))
train_op = optimizer.minimize(loss)

训练和模型保存无需额外 API

saver = tf.compat.v1.train.Saver()
with tf.compat.v1.Session() as sess:
  for i in range(epoch):
    sess.run(train_op)
  saver.save(sess, 'checkpoint/model')

EmbeddingVariable 组件是基于阿里巴巴 PAI 团队 2019 年向 TensorFlow 社区提出 EmbeddingVariable 32 方案实现。

主要目标是在 TensorFlow 框架之下，支持对稀疏参数存储，访问以及更新。

包含以下特性：

1. 动态可伸缩的 EmbeddingVariable (以下简称 EV)，实现特征的动态新增与淘汰

2. 增量 Checkpoint 功能，便于稀疏参数在线训练并实时线上加载模型

3. 支持EV的查询以及相关 Optimizer 的参数更新

4. API 设计兼容原生 TensorFlow，用户侵入性小

5. 支持 Graph Mode 和 Eager Mode

3.2.1 整体设计

使用 HashTable 作为底层数据结构进行存储稀疏参数，目前开源版本使用的是 Google 的 DenseHashMap。封装为 TensorFlow 的 Resource，由 TF 的 ResourceMgr 统一进行管理。Resource 封装在各个与 EV 相关 Operation 中，对 Resource 进行操作（稀疏参数的查询，更新，保存，恢复等）。用户可以直接调用 Operation 或调用封装的 Python API。

tfra中的两种 embedding 稀疏参数的使用方法

3.2.2 底层数据结构

EV 在底层数据结构使用了 HashMap + Tensor 的结构，HashMap 负责管理 KV 元数据，Tensor 存储 Embedding 及版本信息，其中 Tensor 中 version 存储最近一次更新 global_step 数，以便在特征淘汰的时候使用该信息。

tfra中的两种 embedding 稀疏参数的使用方法

image1080×318 13.2 KB

3.2.3 算子实现与 API

围绕 Embedding 的语义，我们实现了对稀疏参数的查询、更新等算子，支持 Graph 和 Eager 模式

EVHandleOp
InitializeEVOp
EVIsInitializedOp
EVShapeOp
EVGatherOp
EVSparseApply{GradientDescent, Adagrad, Adam}
EVImport
EVExport

API 主要分为 EV 变量创建以及 Optimizer

tfra.embedding_variable.EmbeddingVariable(name, embedding_dim, ...) 其他参数同 tf.Variable
tfra.embedding_variable.GradientDescentOptimizer(...) 参数同 TensorFlow
tfra.embedding_variable.AdagradOptimizer(...) 参数同 TensorFlow
tfra.embedding_variable.AdamOptimizer(...) 参数同 TensorFlow

3.2.4 模型保存与恢复

我们采用全量+增量的方式保存稀疏模型

1. 全量模型

通过 EVImport/EVExport 两个 Op 对 EV 进行保存与恢复，自动在tf.train.Saver 中完成构图。

2. 增量模型

稀疏参数在训练过程中往往具有局部性特点，参数更新数量占全部参数比重很小，我们在增量模型保存时候只保存相对上一次的增量模型，这样节省了 I/O，而且使得部署增量模型更加快速，适合在线学习的场景。

具体设计为，在每次参数更新的时候记录更改的参数，统计一定间隔时间内参数更改，最后再调用参数保存。逻辑均通过 Operation在Graph 中表达，并且 checkpoint 数据格式兼容 TensorFlow。

3.2.5 使用示例

使用 HashTable 作为底层数据结构进行存储稀疏参数，目前开源版本使用的是 Google 的 DenseHashMap。封装为 TensorFlow 的 Resource，由 TF 的 ResourceMgr 统一进行管理。Resource 封装在各个与 EV 相关 Operation 中，对 Resource 进行操作（稀疏参数的查询，更新，保存，恢复等）。用户可以直接调用 Operation 或调用封装的 Python API。

加载 TensorFlow 及 TFRA 包

import tensorflow as tf
import tensorflow_recommenders_addons as tfra

创建 EmbeddingVariable 稀疏参数

embedding_var = tfra.embedding_variable.EmbeddingVariable(
    name = "embedding_var",
    embedding_dim = 16,
    initializer = tf.keras.initializers.RandomNormal(-1, 1))

兼容原生 TensorFlow API 的稀疏参数查询

ids = ...
embedding_val = tf.nn.embedding_lookup(
    params = embedding_var,
    ids = ids)

计算模型损失函数

loss = …

创建支持普通 Variable 以及 EV 的优化器

optimizer = tfra.embedding_variable.AdamOptimizer(learning_rate=0.001)
train_op = optimizer.minimize(loss)

创建 Saver

saver = tf.train.Saver()

在 Graph Mode 下迭代训练并保存模型

with tf.compat.v1.Session() as sess:
    for i in range(epoch):
        loss_t, _ = sess.run([loss, train_op])
        print("epoch:", i, "loss:" loss_t)
    saver.save(sess, "checkpoint/model")

我们 2021 年 3 月 25 日在中文社区已经进行了一次社区分享，相关 ppt 和视频大家可以在此处 76 找到（提取码: pqpt）

围绕稀疏参数在业务场景中还有很多功能有待完善，我们也将陆续开源 HashtableOnGPU、常用算子的 GPU 版本、特征过滤 API、AdaptiveEmbedding、Frequency Awareness Embedding 等功能。

TensorFlow Recommenders-Addons 现已在 GitHub 上开源，我们的目标是让其不断发展，能够灵活地进行学术研究以及工业应用，并以高度可扩展的方式构建全网推荐系统。

我们希望此文能让您对 TensorFlow Recommenders Addons 有所了解。更多信息，请查看我们的教程 47 或 API 参考 18。我们已经成立 TensorFlow SIG Recommenders 兴趣小组，一同推动基于 TensorFlow 的开源大规模推荐系统的发展，请考虑贡献您的一份力量 134！欢迎大家就嵌入向量学习和分布式训练与应用等主题开展合作和做出贡献。

正文完

请博主喝杯咖啡吧！