TensorFlow 模型准实时更新上线的设计与实现

一点资讯技术团队
2,257 阅读14分钟

前言

一点机器学习平台主要为信息流推荐系统服务，其中一部分业务的模型训练由深度学习框架 TensorFlow 承担。我们在实践过程中发现，原生 TensorFlow 在用于大规模稀疏模型场景中有一些不足和可以优化的地方，比如

  1. 大规模稀疏模型参数的同步和上线速度较慢，训练效果不能及时应用于推理服务

  2. 单机预测时稀疏特征 Embedding 内存占用较大，难以支持所需的大规模模型参数

针对以上问题，我们基于对 TensorFlow 框架的理解进行了定制优化。

问题分析

模型训练的参数保存

TensorFlow 原有的参数保存方案是在训练过程中，将参数保存为文件存放在磁盘中，供预测时使用。推荐系统中存在大量高维离散稀疏特征，参数在模型训练阶段有以下特点

推荐模型的更新具有稀疏性，即一段时间内训练更新的参数只占总量很少的一部分 对于每个 Embedding 参数来说，参数只在矩阵的部分位置发生梯度变化，更新整个 Embedding 参数存在重复的数据 IO 针对上述特点可以采取以下优化措施

在保存模型参数时，由于只有部分参数发生了变化，那么可以只对发生梯度变化的参数进行保存。相比于前一轮训练迭代没有变化的参数，在增量更新数据时没有重复保存的必要 某一 Embedding 参数在保存时，由于参数矩阵稀疏的特点，梯度变化只会发生在矩阵的局部。同样的，在增量覆盖保存时，不变的部分没有再次保存的必要，只需要更新发生梯度变化的部分 通过增量保存数据的方式，聚集于变化的数据进行处理，在大规模参数稀疏性的特点下，可极大的减小所需同步的数据，所需同步数据量的减小可以带来数据同步效率的提升。

另外，从工程实施角度考虑，对同步数据用批量处理，IO 交互的减少又能进一步的提高同步效率。

模型同步

TensorFlow 原有的模型参数上线流程，需要在训练中将参数保存到文件，再将文件传输到预测服务器，由预测服务器进行加载使用。这个过程流程长，文件传输慢，更新不及时。

如果将模型训练时保存的参数和预测服务共用一套，就可几乎完全节省掉参数同步过程的时间消耗。尤其是对于大规模参数的数据传输，节省同步时间带来的效率提升就更大。

当然出于数据隔离的考虑，这种方式还需要模型参数的版本管理等辅助支持。

模型使用

大规模模型推理的场景中，内存的瓶颈是很大的挑战。分布式训练模式将参数分布在 PS 集群，可以缓解内存中的参数数据的存储压力。但面对高请求 qps 时推理服务缺乏很好的支持方案，即使采用同分布式训练相同的分布式方案，参数在集群间传输会带来较高的网络带宽压力和响应延迟。

如果可以做到精简加载数据，减小预测服务所需的内存要求，就可以较好的解决内存瓶颈的问题。因为部分参数稀疏的特点，可以想到用数据分片的思想，用多个局部数据来表示完整的有效数据。

同时，稀疏参数实际用到的部分往往也只是整个参数的局部，那么可以考虑只将参数的局部数据进行加载使用，从而降低对内存的存储空间要求。

增强定制

经过以上分析，问题和解决思路明了，我们希望可以对 TensorFlow 增强如下功能

在训练阶段参数可及时保存到指定的存储系统，这里我们采用公司级的 KV 数据库 Morpheus 参数保存时只保存变化的部分参数，并且每个 Embedding 参数只保存更新的部分维度 预测过程中只加载参数被用到的有效数据部分 为达成上述目标，需要对 TensorFlow 进行定制。那么如何着手定制呢，我们再来学习下 TensorFlow 的设计

认识 TensorFlow 计算图和操作 OP

1.gif

TensorFlow 用“结点”（nodes）和“边”（edges）的有向图来描述数学计算。“节点” 一般用来表示施加的数学操作，但也可以表示数据输入（feed in）的起点/输出（push out）的终点，或者是读取/写入持久变量（persistent variable）的终点。

“边”表示“节点”之间的输入/输出关系。这些数据“边”可以输运“size可动态调整”的多维数据数组，即“张量”（tensor）。张量从图中流过的直观图像是这个工具取名为“TensorFlow”的原因。一旦输入端的所有张量准备好，节点将被分配到各种计算设备完成异步并行地执行运算。

简言之，计算图结构由模型的算法结构决定，对数据的操作即为 operation（ op ）。当模型结构确定的情况下，我们的增强就需要对 op 进行定制。

那么需要对哪些 op 进行增强呢，我们先结合一个简单的代码示例来观察一个简单的计算图，以便了解一次数据流图计算过程的关键过程。例如

import tensorflow as tf
from tensorflow.core.protobuf import saver_pb2
from tensorflow.python.ops import variables
from numpy import random
 
a_matrix = random.random(size=(2,4))
b_matrix = random.random(size=(2,4))
print("a_matrix=", a_matrix)
print("b_matrix=", b_matrix)
a = tf.Variable(a_matrix, dtype=tf.float32, name="a")
b = tf.Variable(b_matrix, dtype=tf.float32, name="b")
 
res_a = tf.nn.embedding_lookup(a, [0, 0], name="lookup_a")
res_b = tf.nn.embedding_lookup(b, [1, 1], name="lookup_b")
y = tf.add(res_a, res_b)
 
saver = tf.train.Saver(variables._all_saveable_objects(), sharded=True, write_version=saver_pb2.SaverDef.V2,  allow_empty=True)
meta_graph_def = saver.export_meta_graph(as_text=True, clear_devices=True, strip_default_attrs=True)
 
with open("./meta_graph_def.pbtxt", "w") as f:
    f.write(str(meta_graph_def))
 
with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    ## tensorboard --logdir=tensorboard/test0/
    writer = tf.summary.FileWriter("./tensorboard/test0/", sess.graph)
    print("res_a=", sess.run(res_a))
    print("res_b=", sess.run(res_b))
    print("y=", sess.run(y))
    writer.close()
 
if __name__ == '__main__':
    print("hello world")

这个计算图实现了简单的计算，定义两个变量，分别执行 embedding_lookup，然后对查询结果求和计算。代码中对 graph 进行保存，也保存了 tensorboard 所需的数据，用于进一步分析。

使用 tensorboard 打开文件保存的路径，我们可以看到这个计算图的直观表现。分别查看计算节点，就可以观察到 op 的输入输出关系。

2.png

3.png

4.png

定制 OP 与定制计算图

考虑到了官方库里 op 无法涵盖用户需求的情况，TensorFlow 提供了自定义 op 操作的扩展能力。知道了数据定义、保存、采集等操作，那么就可以进行定制化的修改。通过编写自定义的 op，使数据在保存时，按设计好的格式保存到我们指定的数据库。数据在加载时，又从数据库反序列化并组装为所需的参数，从而绕过 TensorFlow 框架从文件加载参数的方式。以下列举部分定制 op 的设计功能以作参考

5.png

定制好 op 后，如何替换模型计算图中原生的 op 呢？TensorFlow 在模型保存时，会生成 meta_graph_def 文件，文件内容是采用类似 json 的格式描述计算图的结构关系。当加载此文件时，TensorFlow 会根据文件中描述的结构信息构建出计算图。

可以修改模型保存的 meta_graph_def 文件，将其中的 op 替换为我们定制的 op，同时修改每个 node 的 input 和 output 关系，以修改 op 之间的依赖关系。随后用修改过的 meta_graph_def 文件加载回模型的计算图，即完成了对原有计算图结构的修改。如下为修改 meta_graph_def 文件的局部示例。

node {
  name: "a/read"
  op: "InferenceIdentity" #### 原始op为原生的Identity，这里将其替换为了定制的op
  input: "a"
  attr {
    key: "T"
    value {
      type: DT_FLOAT
    }
  }
  ...
}

经过回顾编写自定义 op 和对 meta_graph_def 修改的方案，就基本做好了增强定制的准备工作。

参数增量更新

训练过程一般用 minimize 函数反向传播更新参数收敛 loss。minimize 函数内部是由两步完成的

第一步，compute_gradients 函数返回更新的梯度和参数 grads_and_vars

第二步，apply_gradients 函数将参数和梯度应用于现有参数，完成参数的更新

在第一步中，获得本次需要更新的梯度和参数，其中 Embedding 参数的梯度中包含每个 tensor 中发生变化的数据切片 IndexedSlices。关于 IndexedSlices，官方文档中部分描述如下。

“ An IndexedSlices is typically used to represent a subset of a larger tensor dense of shape [LARGE0, D1, .. , DN] where LARGE0 >> D0. The values in indices are the indices in the first dimensIOn of the slices that have been extracted from the larger tensor. ”

可以认为是一种类似 SparseTensor 的思想，用元素数据和元素位置表示一个较大 tensor 。描述中是将 tensor 按第一维度切片，从而将一个较大的形状为 [LARGE0, D1, .. , DN] 的 tensor 表示为多个较小的形状为 [D1, .. , DN] 的 tensor。

6.png

跟据 IndexedSlices，可以获得本次梯度变化对应 tensor 变化的部分。在同步参数时，可以按 IndexedSlices 分好的切片，将需要的切片数据部分保存到数据库中。例如一个形状为 [m, n, l, k] 的 tensor 可以按切片的数量保存为 m 个形状为 [n, l, k] 的 KV 数据，key 为 tensor_name 和 m 维度序号组合的唯一命名。

通过将 optimizer 分为两步，在每轮训练的迭代过程中，RecordIndicesOp 记录下梯度更新的参数及它们梯度 IndexedSlices 的 indices。那么就可以持续追踪训练更新的参数和参数发生变化的部分。

经过多轮训练迭代后，MorpheusWriteOp 将追踪记录的映射数据按 indices 写出到 KV 存储，就将这段时间内的变化参数同步到数据库。参数写出完成后 ClearIndicesOp 清除记录的 index 数据，即可开始下一轮增量更新的周期。

如下图示，从 m 轮训练迭代开始，稀疏参数 a 在局部被更新数值，recorder 记录下被更新的第一维度的坐标。经过 n 轮迭代后，结合当前 a 参数和 record 中记录的变化过的第一维度坐标，取出参数发生变化的部分，将数据写出到 KV 数据库中。随后清除掉 recorder 的记录，开始下一轮的迭代。

7.png

模型的加载

根据前面的过程，我们已经得到定制化后的模型 graph，和已经保存在 KV 数据库中的参数数据。在推理服务中，计算图 metra_graph 由原生方式加载。参数部分的加载由前面提到的定制 op 处理，这里展开介绍。

在加载完计算图结构后，指定一个 fake 文件路径加载参数数据。需要注意的是前面我们已经修改了用于加载数据的 RestoreV2 op ，所以在此处加载时并没有任何参数被加载，加载 fake 文件只是为了满足框架流程进行初始化工作。

那么保存在 KV 数据库的参数数据是如何加载的呢？TensorFlow 的 graph 在执行时会自动运行所需的依赖，当 session 在第一次 run 时数据会流经定制 op ，从而触发定制部分的操作。所以在加载完模型图结构和 fake 的参数文件后，还需要再执行一次 inference 完成对 fake 参数的替换。替换过程是使用 InferenceIdentityOp 从数据库中获取数据并按 shape 组装好需要的 tensor， 换下之前的 fake tensor 并放进模型。对于 Embedding 数据加载又有一些特殊优化，在下一部分会做讲解。

推理服务的数据使用

对预测服务来说，一般在 inference 前需要完成模型和数据的加载。非 Embedding 数据一般较小，不会构成工程实践的挑战。但 Embedding 数据大小与词表和维度相关， [vocabulary_size, Embedding_dimension] 往往很大，对于推荐业务来说，特征层参数又有稀疏的特点。那么加载完整的 Embedding 向量从内存角度说是很浪费的。而且当词表很大时，大规模参数对于单机内存很难承受无法加载，这就带来较大的工程挑战。

针对大规模 Embedding 参数的定制处理

实际在使用中 Embedding 参数会按 embedding_lookup 从完整的参数中查找部分数据 ，embedding_sparse_lookup 同理。那么可以只加载完整参数需要使用的某些部分，而不必加载整个 Embedding 参数。

在 “参数增量更新” 章节中介绍到对于梯度变化的参数保存方案，Embedding 数据作为训练中迭代参数的一个子集，也用同样的数据切片方案保存了下来。在 “ 定制op ” 章节中提到，在预测服务中，实际是使用定制的 InferenceGatherV2Op 来进行 embedding_lookup 操作。以下举例来讲解 Embedding 数据从训练到推理的完整流程。

例如，在训练中定义

a = tf.Variable(a_matrix, dtype=tf.float32, name="a")

在 inference 中使用

b = tf.nn.embedding_lookup(a, [0, 1, 3], name="b")

参数 a 和参数 b 的流转过程如下图表示

8.png

在训练过程中，参数 a（shape [4,4]）发生梯度变化，经过 InferenceRecordOp 的处理，按梯度变化 IndexedSlices 切片后，将数据存储在 KV 数据库中。参数 a 行号为 0 和 2 的部分，分别被保存为 “a:0” = [1, 0, 0, 0] 和 “a:2" = [0, 1, 0, 0] 的两条 KV 数据，其余行号为 1 和 3 的部分并未保存在数据库中。

在模型上线初始化过程中，并未将参数 a 加载，而是加载了一条 fake 数据 a' ( shape [1] ) 以保证模型图可以正常加载。

当 inference 执行过程中需要对参数 a 进行 embedding_lookup ( a, [0, 1, 3] ) 时，InferenceGatherV2Op 会尝试从 KV 数据库中加载 key 为 “a:0”、“a:1”、“a:3” 的 3 条数据。实际只可以获取到 “a:0” 这一条数据，然后根据组装规则，将获取的数据组装为数据 b（shape [3,4]），完成 embedding_lookup 操作。

可以看出整个过程参数 a 并未完整的加载在预测服务中，但并没有影响推理过程对参数 a 的使用。

数据同步的优化

以上设计在预测过程中由 KV 数据库充当参数服务器，使用公司的 KV 系统 Morpheus 来做底层支撑，减少了数据同步过程，提高了训练数据上线的效率。为了降低对 KV 数据库的读写压力，Embedding 之外的参数数据在预测服务中做了本地缓存。但因为本地缓存的存在，客户端拉取数据请求会被缓存屏蔽，无法及时获取最新参数。所以针对这个情况，Morpheus 客户端使用订阅消息队列的方式，提供了监听数据更新的功能。当 Morpheus server 数据更新时，client 也会感知到变化的 key，然后拉取数据更新至本地缓存，做了近乎实时的参数更新。

总结与展望

以上主要根据数据分片和数据懒加载的思想，尝试从参数同步和数据使用的方面对 TensorFlow 框架进行增强定制，实现了准实时的模型训练参数更新。相比原有的小时级模型参数更新，推荐系统指标取得了显著的提升。

但随着推荐算法对数据维度、词表大小、应用效率等要求的不断提升，我们仍然面对不少挑战。我们会持续基于 TensorFlow 的工程应用做更多的探索和尝试。如果您对 TensorFlow 也感兴趣，欢迎大家在文末评论区留言或者给出建议，非常感谢。

文章来自一点资讯机器学习平台团队

avatar
文章
阅读
粉丝