深度学习框架TensorFlow在Kubernetes上的实践

torlee

什么是TensorFlow
　　TensorFlow是谷歌在去年11月份开源出来的深度学习框架。开篇我们提到过AlphaGo，它的开发团队DeepMind已经宣布之后的所有系统都将基于TensorFlow来实现。TensorFlow一款非常强大的开源深度学习开源工具。它可以支持手机端、CPU、GPU以及分布式集群。TensorFlow在学术界和工业界的应用都非常广泛。在工业界，基于TensorFlow开发的谷歌翻译、谷歌RankBrain等系统都已经上线。在学术界很多我在CMU、北大的同学都表示TensorFlow是他们实现深度学习算法的首选工具。

　　上面的ppt给出了一个简单的TensorFlow程序样例，这个样例实现了向量加法的功能。TensorFlow提供了Python和C++的API，但Python的API更全面，所以大部分TensorFlow程序都是通过Python实现的。在上面程序的第一行我们通过import将TensorFlow加载进来。在TensorFlow中所有的数据都是通过张量（Tensor）的方式存储，要计算张量中数据的具体取值，我们需要通过一个会话（session）。
　　上面代码中的第二行展示了如何生成会话。会话管理运行一个TensorFlow程序所需要的计算资源。TensorFlow中一个比较特殊的张量是变量（tf.Variable），在使用变量之前，我们需要明确调用变量初始化的过程。在上面的代码最后一行，我们可以看到要得到结果张量output的取值，我们需要明确调用计算张量取值的过程。

　　通过TensorFlow实现神经网络是非常简单的。通过TFLearn或者TensorFlow-Slim可以在10行之内实现MNIST手写体数字识别问题。上面的ppt展示了TensorFlow对于不同神经网络结构的支持，可以看出，TensorFlow可以在很短的代码内支持各种主要的神经网络结构。
　　虽然TensorFlow可以很快的实现神经网络的功能，不过单机版的TensorFlow却很难训练大规模的深层神经网络。

　　这张图给出了谷歌在2015年提出的Inception-v3模型。这个模型在ImageNet数据集上可以达到95%的正确率。然而，这个模型中有2500万个参数，分类一张图片需要50亿次加法或者乘法运算。即使只是使用这样大规模的神经网络已经需要非常大的计算量了，如果需要训练深层神经网络，那么需要更大的计算量。神经网络的优化比较复杂，没有直接的数学方法求解，需要反复迭代。在单机上要把Inception-v3模型训练到78%的准确率大概需要5个多月的时间。如果要训练到95%的正确率需要数年。这对于实际的生产环境是完全无法忍受的。
TensorFlow on Kubernetes
　　如我们上面所介绍的，在单机环境下是无法训练大型的神经网络的。在谷歌的内部，Google Brain以及TensorFlow都跑在谷歌内部的集群管理系统Borg上。我在谷歌电商时，我们使用的商品分类算法就跑在1千多台服务器上。在谷歌外，我们可以将TensorFlow跑在Kubernetes上。在介绍如何将TensorFlow跑在Kubernetes上之前，我们先来介绍一下如何并行化的训练深度学习的模型。

　　深度学习模型常用的有两种分布式训练方式。一种是同步更新，另一种是异步更新。如上面的ppt所示，在同步更新模式下，所有服务器都会统一读取参数的取值，计算参数梯度，最后再统一更新。而在异步更新模式下，不同服务器会自己读取参数，计算梯度并更新参数，而不需要与其他服务器同步。同步更新的最大问题在于，不同服务器需要同步完成所有操作，于是快的服务器需要等待慢的服务器，资源利用率会相对低一些。而异步模式可能会使用陈旧的梯度更新参数导致训练的效果受到影响。不同的更新模式各有优缺点，很难统一的说哪一个更好，需要具体问题具体分析。
　　无论使用哪种更新方式，使用分布式TensorFlow训练深度学习模型需要有两种类型的服务器，一种是参数服务器，一种是计算服务器。参数服务器管理并保存神经网络参数的取值；计算服务器负责计算参数的梯度。
　　在TensorFlow中启动分布式深度学习模型训练任务也有两种模式。一种为In-graph replication。在这种模式下神经网络的参数会都保存在同一个TensorFlow计算图中，只有计算会分配到不同计算服务器。另一种为Between-graph replication，这种模式下所有的计算服务器也会创建参数，但参数会通过统一的方式分配到参数服务器。因为In-graph replication处理海量数据的能力稍弱，所以Between-graph replication是一个更加常用的模式。

　　最后一个问题，我们刚刚提到TensorFlow是支持以分布式集群的方式运行的，那么为什么还需要Kubernetes？如果我们将TensorFlow和Hadoop系统做一个简单的类比就可以很清楚的解释这个问题。大家都知道Hadoop系统主要可以分为Yarn、HDFS和mapreduce计算框架，那么TensorFlow就相当于只是Hadoop系统中Mapreduce计算框架的部分。
　　TensorFlow没有类似Yarn的调度系统，也没有类似HDFS的存储系统。这就是Kubernetes需要解决的部分。Kubernetes可以提供任务调度、监控、失败重启等功能。没有这些功能，我们很难手工的去每一台机器上启动TensorFlow服务器并时时监控任务运行的状态。除此之外，分布式TensorFlow目前不支持生命周期管理，结束的训练进程并不会自动关闭，这也需要进行额外的处理。