深入浅出TensroFlow(七)TensorFlow计算加速

作者 | 郑泽宇

AI前线出品| ID:ai-front

2017年2月16日,Google正式对外发布Google TensorFlow 1.0版本,并保证本次的发布版本API接口完全满足生产环境稳定性要求。这是TensorFlow的一个重要里程碑,标志着它可以正式在生产环境放心使用。在国内,从InfoQ的判断来看,TensorFlow仍处于创新传播曲线的创新者使用阶段,大部分人对于TensorFlow还缺乏了解,社区也缺少帮助落地和使用的中文资料。InfoQ期望通过深入浅出TensorFlow系列文章能够推动Tensorflow在国内的发展。欢迎加入QQ群(群号:183248479)深入讨论和交流。下面为本系列的前六篇文章:

深入浅出Tensorflow(一):深度学习及TensorFlow简介

深入浅出TensorFlow(二):TensorFlow解决MNIST问题入门

深入浅出Tensorflow(三):训练神经网络模型的常用方法

深入浅出Tensorflow(四):卷积神经网络

深入浅出Tensorflow(五):循环神经网络简介

深入浅出TensorFlow(六)TensorFlow高层封装

在前面的文章中介绍了使用TensorFlow实现各种深度学习的算法。然而要将深度学习应用到实际问题中,一个非常大的问题在于训练深度学习模型需要的计算量太大。比如要将Inception-v3模型在单机单卡上训练到78%的正确率需要将近半年的时间,这样的训练速度是完全无法应用到实际生产中的。为了加速训练过程,本文将介绍如何通过TensorFlow利用GPU或/和分布式计算进行模型训练。

TensorFlow使用GPU

TensorFlow程序可以通过tf.device函数来指定运行每一个操作的设备,这个设备可以是本地的CPU或者GPU,也可以是某一台远程的服务器。TensorFlow会给每一个可用的设备一个名称,tf.device函数可以通过设备的名称来指定执行运算的设备。比如CPU在TensorFlow中的名称为/cpu:0。

在默认情况下,即使机器有多个CPU,TensorFlow也不会区分它们,所有的CPU都使用/cpu:0作为名称。而一台机器上不同GPU的名称是不同的,第n个GPU在TensorFlow中的名称为/gpu:n。比如第一个GPU的名称为/gpu:0,第二个GPU名称为/gpu:1,以此类推。

TensorFlow提供了一个快捷的方式来查看运行每一个运算的设备。在生成会话时,可以通过设置log_device_placement参数来打印运行每一个运算的设备。下面的程序展示了如何使用log_device_placement这个参数。

深入浅出TensroFlow(七)TensorFlow计算加速_第1张图片

在以上代码中,TensorFlow程序生成会话时加入了参数log_device_placement=True,所以程序会将运行每一个操作的设备输出到屏幕。于是除了可以看到最后的计算结果之外,还可以看到类似“add:/job:localhost/replica:0/task:0/cpu:0”这样的输出。这些输出显示了执行每一个运算的设备。比如加法操作add是通过CPU来运行的,因为它的设备名称中包含了/cpu:0。

在配置好GPU环境的TensorFlow中,如果操作没有明确地指定运行设备,那么TensorFlow会优先选择GPU。比如将以上代码在亚马逊(Amazon Web Services, AWS)的 g2.8xlarge实例上运行时,会得到以下运行结果。

深入浅出TensroFlow(七)TensorFlow计算加速_第2张图片

从上面的输出可以看到在配置好GPU环境的TensorFlow中,TensorFlow会自动优先将运算放置在GPU上。不过,尽管g2.8xlarge实例有4个GPU,在默认情况下,TensorFlow只会将运算优先放到/gpu:0上。于是可以看见在上面的程序中,所有的运算都被放在了/gpu:0上。如果需要将某些运算放到不同的GPU或者CPU上,就需要通过tf.device来手工指定。下面的程序给出了一个通过tf.device手工指定运行设备的样例。

深入浅出TensroFlow(七)TensorFlow计算加速_第3张图片

在以上代码中可以看到生成常量a和b的操作被加载到了CPU上,而加法操作被放到了第二个GPU“/gpu:1”上。在TensorFlow中,不是所有的操作都可以被放在GPU上,如果强行将无法放在GPU上的操作指定到GPU上,那么程序将会报错。

多GPU并行

下面将给出具体的TensorFlow代码在一台机器的多个GPU上并行训练深度学习模型。因为一般来说一台机器上的多个GPU性能相似,所以在这种设置下会更多地采用同步模式训练深度学习模型。下面将给出具体的代码,在多GPU上训练深度学习模型解决MNIST问题。(该代码可以在TensorFlow 0.9.0下运行,对于更新TensorFlow的版本,请参考Github代码库:https://github.com/caicloud/tensorflow-tutorial)

深入浅出TensroFlow(七)TensorFlow计算加速_第4张图片
深入浅出TensroFlow(七)TensorFlow计算加速_第5张图片
深入浅出TensroFlow(七)TensorFlow计算加速_第6张图片
深入浅出TensroFlow(七)TensorFlow计算加速_第7张图片
深入浅出TensroFlow(七)TensorFlow计算加速_第8张图片
深入浅出TensroFlow(七)TensorFlow计算加速_第9张图片
深入浅出TensroFlow(七)TensorFlow计算加速_第10张图片
深入浅出TensroFlow(七)TensorFlow计算加速_第11张图片
深入浅出TensroFlow(七)TensorFlow计算加速_第12张图片
深入浅出TensroFlow(七)TensorFlow计算加速_第13张图片

图1 在AWS的g2.8xlarge实例上运行MNIST样例程序时GPU的使用情况

在AWS的g2.8xlarge实例上运行以上代码可以同时使用4个GPU训练神经网络。图1显示了运行样例代码时不同GPU的使用情况。因为运行的神经网络规模比较小,所以在图1中显示的GPU使用率不高。如果训练大型的神经网络模型,TensorFlow将会占满所有用到的GPU。

深入浅出TensroFlow(七)TensorFlow计算加速_第14张图片

图2 使用了4个GPU的TensorFlow计算图可视化结果

图2展示了通过TensorBoard可视化得到的样例代码TensorFlow计算图,其中节点上的颜色代表了不同的设备,比如黑色代表CPU、白色代表第一个GPU,等等。从图2中可以看出,训练神经网络的主要过程被放到了GPU_0、GPU_1、GPU_2和GPU_3这4个模块中,而且每一个模块运行在一个GPU上。

分布式TensorFlow

通过多GPU并行的方式可以达到很好的加速效果。然而一台机器上能够安装的GPU有限,要进一步提升深度学习模型的训练速度,就需要将TensorFlow分布式运行在多台机器上。以下代码展示了如何创建一个最简单的TensorFlow集群:

深入浅出TensroFlow(七)TensorFlow计算加速_第15张图片

在以上代码中,首先通过 tf.train.Server.create_local_server函数在本地建立了一个只有一台机器的TensorFlow集群。然后在该集群上生成了一个会话,并通过生成的会话将运算运行在创建的TensorFlow集群上。虽然这只是一个单机集群,但它大致反应了TensorFlow集群的工作流程。TensorFlow集群通过一系列的任务(tasks)来执行TensorFlow计算图中的运算。一般来说,不同任务跑在不同机器上。最主要的例外是使用GPU时,不同任务可以使用同一台机器上的不同GPU。TensorFlow集群中的任务也会被聚合成工作(jobs),每个工作可以包含一个或者多个任务。比如在训练深度学习模型时,一台运行反向传播的机器是一个任务,而所有运行反向传播机器的集合是一种工作。

上面的样例代码是只有一个任务的集群。当一个TensorFlow集群有多个任务时,需要使用tf.train.ClusterSpec来指定运行每一个任务的机器。比如以下代码展示了在本地运行有两个任务的TensorFlow集群。第一个任务的代码如下:

深入浅出TensroFlow(七)TensorFlow计算加速_第16张图片

下面给出了第二个任务的代码:

深入浅出TensroFlow(七)TensorFlow计算加速_第17张图片

启动第一个任务后,可以得到类似下面的输出:

深入浅出TensroFlow(七)TensorFlow计算加速_第18张图片

从第一个任务的输出中可以看到,当只启动第一个任务时,程序会停下来等待第二个任务启动,而且持续输出failed to connect to 'ipv4:127.0.0.1:2223': socket error: connection refused。当第二个任务启动后,可以看到从第一个任务中会输出Hello from server1!的结果。第二个任务的输出如下:

深入浅出TensroFlow(七)TensorFlow计算加速_第19张图片

值得注意的是第二个任务中定义的计算也被放在了设备/job:local/replica:0/task:0/cpu:0上。也就是说这个计算将由第一个任务来执行。从上面这个样例可以看到,通过tf.train.Server.target生成的会话可以统一管理整个TensorFlow集群中的资源。

和使用多GPU类似,TensorFlow支持通过tf.device来指定操作运行在哪个任务上。比如将第二个任务中定义计算的语句改为以下代码,就可以看到这个计算将被调度到/job:local/replica:0/task:1/cpu:0上面。

在上面的样例中只定义了一个工作“local”。但一般在训练深度学习模型时,会定义两个工作。一个工作专门负责存储、获取以及更新变量的取值,这个工作所包含的任务统称为参数服务器(parameter server,ps)。另外一个工作负责运行反向传播算法来获取参数梯度,这个工作所包含的任务统称为计算服务器(worker)。下面给出了一个比较常见的用于训练深度学习模型的TensorFlow集群配置方法。

深入浅出TensroFlow(七)TensorFlow计算加速_第20张图片

使用分布式TensorFlow训练深度学习模型一般有两种方式。一种方式叫做计算图内分布式(in-graph replication)。使用这种分布式训练方式时,所有的任务都会使用一个TensorFlow计算图中的变量(也就是深度学习模型中的参数),而只是将计算部分发布到不同的计算服务器上。

上面给出的使用多GPU样例程序就是这种方式。多GPU样例程序将计算复制了多份,每一份放到一个GPU上进行运算。但不同的GPU使用的参数都是在一个TensorFlow计算图中的。因为参数都是存在同一个计算图中,所以同步更新参数比较容易控制。在上面给出的代码也实现了参数的同步更新。然而因为计算图内分布式需要有一个中心节点来生成这个计算图并分配计算任务,所以当数据量太大时,这个中心节点容易造成性能瓶颈。

另外一种分布式TensorFlow训练深度学习模型的方式叫计算图之间分布式(between-graph replication)。使用这种分布式方式时,在每一个计算服务器上都会创建一个独立的TensorFlow计算图,但不同计算图中的相同参数需要以一种固定的方式放到同一个参数服务器上。TensorFlow提供了tf.train.replica_device_setter函数来帮助完成这一个过程。

因为每个计算服务器的TensorFlow计算图是独立的,所以这种方式的并行度要更高。但在计算图之间分布式下进行参数的同步更新比较困难。

为了解决这个问题,TensorFlow提供了tf.train.SyncReplicasOptimizer函数来帮助实现参数的同步更新。这让计算图之间分布式方式被更加广泛地使用。因为篇幅所限,所以这里不在给出具体代码,具体代码可以在Github代码库中找到:

https://github.com/caicloud/tensorflow-tutorial/tree/master/Deep_Learning_with_TensorFlow/1.0.0/Chapter10

使用Caicloud TaaS平台运行分布式TensorFlow

每次运行分布式TensorFlow都需要登录不同的机器来启动集群。这使得使用起来非常不方便。当需要使用100台机器运行分布式TensorFlow时,需要手动登录到每一台机器并启动TensorFlow服务,这个过程十分繁琐。而且,当某个服务器上的程序死掉之后,TensorFlow并不能自动重启,这给监控工作带来了巨大的难度。

如果类比TensorFlow与Hadoop,可以发现TensorFlow只实现了相当于Hadoop中MapReduce的计算框架,而没有提供类似Yarn的集群管理工具以及HDFS的存储系统。为了降低分布式TensorFlow的使用门槛,才云科技(Caicloud.io)基于Kubernetes容器云平台提供了一个分布式TensorFlow平台TensorFlow as a Service(TaaS)。

从我们提供的开源代码库中可以看出,编写分布式TensorFlow程序需要指定很多与模型训练无关的代码来完成TensorFlow集群的设置工作。为了降低分布式TensorFlow的学习成本,Caicloud的TensorFlow as a Service(TaaS)平台首先对TensorFlow集群进行了更高层的封装,屏蔽了其中与模型训练无关的底层细节。

其次,TaaS平台结合了谷歌开源的容器云平台管理工具Kubernetes来实现对分布式TensorFlow任务的管理和监控,并支持通过UI设置分布式 TensorFlow任务的节点个数、是否使用GPU等信息。

有关Caicloud TaaS 平台的更多信息可以了解

https://caicloud.io/products/taas。

有关TaaS平台的使用文档可以参考

https://docs.caicloud.io/clever/index.html。

本文内容来自作者图书作品《TensorFlow:实战Google深度学习框架》。

作者介绍

郑泽宇,才云首席大数据科学家,前谷歌高级工程师。从 2013 年加入谷歌至今,郑泽宇作为主要技术人员参与并领导了多个大数据项目,拥有丰富机器学习、数据挖掘工业界及科研项目经验。2014 年,他提出产品聚类项目用于衔接谷歌购物和谷歌知识图谱(Knowledge Graph)数据,使得知识卡片形式的广告逐步取代传统的产品列表广告,开启了谷歌购物广告在搜索页面投递的新纪元。他于2013 年 5 月获得美国 Carnegie Mellon University(CMU)大学计算机硕士学位, 期间在顶级国际学术会议上发表数篇学术论文,并获得西贝尔奖学金。


-全文完-

关注人工智能的落地实践,与企业一起探寻 AI 的边界,AICon 全球人工智能技术大会火热售票中,8 折倒计时一周抢票,详情点击:

http://t.cn/Rl2MGtT

《深入浅出TensorFlow》迷你书现已发布,关注公众号“AI前线”,ID:ai-front,回复关键字:TF,获取下载链接!

深入浅出TensroFlow(七)TensorFlow计算加速_第21张图片

你可能感兴趣的:(深入浅出TensroFlow(七)TensorFlow计算加速)