TensorFlow 是目前使用最广泛的机器学习框架之一,它加快了研究速度,并减少了新模型的生产时间。在一篇论文中,来自原生程序化 DSP 公司 Zemanta 的数据科学总监 Davorin Kopič和工程师 Jan Hartman 展示了将在线广告生态系统中的大规模机器学习模型转换为 TensorFlow 框架的过程,并将在 TensorFlow 框架中实现的机器学习模型扩展到每秒超过 3 亿次预测。因此,该研究的主要内容是在 TF 中实现模型并使用各种优化技术以低延迟有效地为其提供服务。
论文地址:https://arxiv.org/abs/2109.09541
该研究使用的案例是线上广告的点击预测。在 RTB (实时竞价)中,多个 DSP(竞标者)通过在网页的加载过程中实时竞标来竞争在线广告空间。广告空间是按广告印象出售的,这使得以市场价值出售虚拟广告空间成为可能。通过使用机器学习,RTB 还使广告商能够最大化其 KPI,例如点击率 (CTR)。估算广告的点击率是 RTB 的核心问题之一,拥有一个好的点击预测模型非常重要。
在 Golang 中实现的基于自定义逻辑回归和分解机 (FM) 的模型,其表达能力受限,并且需要手动实现所有学习程序,这些都会减慢实验速度并限制模型的预测性能。因此,研究者决定采用 TensorFlow 框架,并用表达能力更强的模型替换现有模型。
面临的挑战
基于广告竞标这一具体用例,该研究遇到了一些挑战,共分为实现、服务和优化三个方面。
一方面,为每台机器配备一个或多个顶级 GPU 的成本会高得令人望而却步;另一方面,只有一小群 GPU 的机器将迫使研究过渡到基于服务的架构。鉴于这两种选择都不是特别可取,而且该研究的模型与其他深度学习领域(例如计算机视觉或自然语言处理)的 SOTA 模型相比相对较小,因此该研究不在生产中使用 GPU 进行推断。并且由于该研究的模型使用稀疏权重,其用例也不适合 GPU 工作负载。
实现
为了给 TF 模型实现有效的训练循环,该研究实现和测试了多种方法。高吞吐量在线训练和在 TF 中服务的案例研究很少,而且文档往往不够具体,这迫使研究者必须通读源代码和基准原型才能发现实现过程中的陷阱。
TF 提供了一个庞大的生态系统和大量具有 SOTA 算法实现的库。选择功能丰富的已有实现非常容易,但是研究者发现这些实现大多未经优化,因此他们决定自己实现算法。TF 具有不同抽象级别的 API,不过,一些 API 虽然易于使用,但通常是效率低下的最底层操作(例如 Estimator API)。研究者最终选择了 Keras3,因为它是一个围绕底层 TF 操作的瘦包装器(thin wrapper),并具备高水平的性能,且易于理解。由于 TF 是一个功能和资源都非常丰富的库,因此该研究还必须考虑要在其中实现多少机器学习 pipeline。研究者选择暂时搁置特征转换和交互,仅实现学习算法——尽管它们是最小的可替换部分,但却具备最大的改进潜力。
由于 Golang TF 包装器仅支持预测,因此必须在 Python 中实现训练循环。脚本通过将其标准输入作为子进程实现与 Golang 数据 pipeline 的连接。数据以高效的二进制格式发送而无需解析,与 CSV 格式相比,该方法的速度提高了 25%。然后在后台线程中读取数据,以防止模型在等待数据时空闲。基于此,该研究实现了在整个训练 pipeline 中保持高吞吐量。事实证明,高效的输入和输出也是低延迟预测的关键,该研究通过将所有输入特征连接到单个张量(tensor)中,显著减少了在序列化和复制输入数据上花费的时间。
服务
研究者发现,由于计算密集型神经网络的存在,在使用 Golang TF 装饰器的情况下,DeepFM 模型的 CPU 使用率要高得多。尽管带来了指标的显著提升,但将这种方法扩展到 100% 的流量会带来大量的硬件成本。由于当前全球面临芯片短缺的问题,这意味代价是困难和昂贵的。
显然,降低计算成本是很有必要的。然而缩小神经网络模型规模的同时,也会降低模型的预测性能。在深入研究了 TF 之后,研究者发现如果在计算批处理时增加示例的数量,计算效率会大大提升。这种低线性增长是由于 TF 代码被高度向量化了,TF 也会产生每次计算调用的开销,然后将其分批摊销。考虑到这一点,如果希望减少计算调用的数量,就需要将许多请求连接到一个计算中。
研究者构建了全部包含在一个运行的 bidder 实例的自动批处理系统,以避免网络调用。由于每个实例每秒钟接收数千个传入请求,因此可以保证连接来自众多请求的计算,创建更大的批次。研究者通过一些批处理线程实现了这一点,这些线程接收来自传入请求的数据,创建批处理并在批处理完成后初始化计算。计算过程中,每隔几毫秒即初始化一次以避免超时,因为批处理可能在这个时间窗口中没有填充。这种实现是高度优化的,将计算调用的数量减少到五分之一,同时将 TF 计算的 CPU 占用量减半。
虽然在批处理器线程没有获得 CPU 时间的极少数情况下,会发生请求超时,但只有不足 0.01% 的请求会出现这种情况。研究者观察到平均延迟略有增加(平均约 5 毫秒),流量高峰时可能会更多一些。因此他们实施了 SLA(服务等级协议)和适当的监控手段,以确保延迟的稳定性。鉴于没有大幅增加超时的百分比,这些方法仍是非常有效的,也是这一 TF 服务机制的核心。
优化
研究者在 TF 中实施的模型最初比定制的 FMs 慢得多,为了寻找加速空间,研究者大量使用内置的 TF 分析器来寻找执行时间最长的操作,并尽可能进行了改进。最常见的是各种冗余的 reshape 或转换运算。其中一个更有趣的发现是 Adam 优化器比 Adagrad 慢得多 (大约 50%),尽管二者运算数量上的差异很小。分析器显示,对稀疏权值进行梯度更新需要大量计算时间。这是因为模型的权重是稀疏的 (特征大部分是分类的,因此非常稀疏) ,而优化器没有考虑到这个事实。
由于用 Adagrad 替换 Adam 意味着深度模型性能的显著降低,研究者也寻找了其他解决方案:切换到 Lazy Adam 的优化器被证明是非常有效的,因为它可以非常有效地处理稀疏权重问题。结果显示,其整体加快了超过 40% 的训练速度,与 Adagrad 相接近。
由于使用了自适应优化器(比如 Adam),这也需要存储权重矩和方差,每个参数将存储三个值,将保存的模型大小增加了三倍。然而,这些值实际上并不用于预测,只用于训练。研究者利用这一点构建了优化过程,去掉了这些值的模型,减少了 66% 的数据量,并降低了内存使用量和成本。