亚马逊AWS官方博客

如何在AWS上安装使用分布式TensorFlow

前言

近几年来,深度学习得到了长足的发展,伴随着深度学习的发展,深度学习框架也变得越来越为人们所熟知。TensorFlow是谷歌开源的深度学习框架,自开源以来得到了业界的认可和支持,特别是在TensorFlow分布式版本发布以后,更多的深度学习专业人士开始使用TensorFlow进行分布式的深度学习研究。

深度学习框架底层需要硬件资源的支持,特别需要GPU实例的支持,而AWS云端弹性的GPU实例无疑是深度学习从业人员手中最灵活的资源。

本文将介绍在AWS上使用GPU实例安装配置分布式TensorFlow的过程,希望可以让读者快速搭建自己的深度学习环境,尽快深入到自己的生产应用中或者是研究领域中。

环境准备

首先我们需要为TensorFlow安装准备基础环境,其中包括AWS账号的创建,IAM用户的创建,VPC的划定等工作。有关这些具体工作的细节在本文就不详细讲述了,这些方面的细节请参考相关博文或者技术文档。

准备好账号之后就需要启动两台GPU实例进行设置,考虑到启动实例后需要进行一些软件部署,建议先启动一台GPU实例,安装设置好TensorFlow之后创建实例的AMI镜像,然后通过镜像启动第二台GPU实例,这样比较节省时间。

本文以Ubuntu作为基础环境,所以在启动实例的时候选择操作系统时选择Ubuntu镜像,本例选择的是。

进一步需要做的是选择实例类型,在AWS上的GPU实例有G2和P2两种大的类型。

P2使用了NVIDIA的K80 GPU,实例的具体配置如下:

G2使用了NVIDIA的K520 GPU,实例具体配置如下:


选择你希望使用的实例类型,然后按照EC2启动向导启动该实例。关于EC2启动的不同选项请参考相关文档,这里需要留意的是“置放组”选项,如果我们启动多个EC2实例运行TensorFlow并进行分布式计算,把这些实例放在一个“置放组”内会有效提高实例间的网络通讯效率。

实例启动后通过ssh工具连接到该实例上开始安装过程。

安装TensorFlow

准备好EC2实例后,通过ssh工具连接到实例上,开始以下安装工作。

因为TensorFlow安装需要较长时间,所以建议连接到EC2实例上以后通过screen命令或者tmux命令启动session管理,这样安装过程中出现ssh连接中断也可以继续完成安装工作。

首先需要安装相关的依赖包,具体命令如下:

sudo apt-get update
sudo apt-get upgrade
sudo apt-get install -y build-essential git python-pip libfreetype6-dev libxft-dev libncurses-dev libopenblas-dev gfortran python-matplotlib libblas-dev liblapack-dev libatlas-base-dev python-dev python-pydot linux-headers-generic linux-image-extra-virtual unzip python-numpy swig python-pandas python-sklearn unzip wget pkg-config zip g++ zlib1g-dev
sudo pip install -U pip

安装完依赖包以后接着安装CUDA包,本例中我是直接下载了CUDA 8.0的本地deb包进行安装,安装命令如下,读者在参考使用以下命令时注意修改对应的URL链接和版本号:

wget https://developer.nvidia.com/compute/cuda/8.0/prod/local_installers/cuda-repo-ubuntu1604-8-0-local_8.0.44-1_amd64-deb -O cuda-repo-ubuntu1604-8-0-local_8.0.44-1_amd64.deb
 
sudo dpkg -i cuda-repo-ubuntu1604-8-0-local_8.0.44-1_amd64.deb
sudo apt-get update
sudo apt-get install cuda
CUDA安装比较简单,主要时间消耗在下载CUDA安装包上。安装了CUDA包以后,可以通过nvidia-smi命令查看本地的GPU设备情况,一方面检查GPU硬件情况,一方面也可以确认CUDA安装是否成功。

在我的G2.xlarge实例上安装CUDA并执行nvidia-smi命令后输出如下:

+-----------------------------------------------------------------------------+
| NVIDIA-SMI 367.57                 Driver Version: 367.57                    |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|===============================+======================+======================|
|   0  GRID K520           Off  | 0000:00:03.0     Off |                  N/A |
| N/A   33C    P0     1W / 125W |      0MiB /  4036MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+
 
+-----------------------------------------------------------------------------+
| Processes:                                                       GPU Memory |
|  GPU       PID  Type  Process name                               Usage      |
|=============================================================================|
|  No running processes found                                                 |
+-----------------------------------------------------------------------------+
安装了CUDA后接着安装CuDNN,安装CuDNN需要到NVIDIA网站上注册,注册以后才能下载。这里略去注册过程,读者可以自行到以下网站注册:

https://developer.nvidia.com/rdp/cudnn-download

我下载的是cudnn-8.0-linux-x64-v5.1.tgz文件,下载后的安装命令如下:

tar -zxf ./cudnn-8.0-linux-x64-v5.1.tgz
 
sudo cp -R cuda/lib64/* /usr/local/cuda-8.0/lib64
 
sudo cp cuda/include/cudnn.h /usr/local/cuda-8.0/include/

然后我们需要设置一些环境变量,我将以下配置加入~/.bashrc文件中,并用source ~/.bashrc命令让配置生效。

export CUDA_HOME=/usr/local/cuda-8.0
export CUDA_ROOT=/usr/local/cuda-8.0
export PATH=$PATH:$CUDA_ROOT/bin
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$CUDA_ROOT/lib64

安装了NVIDIA的相关组件后我们还需要安装bazel,为了安装bazel,我们需要先安装java8,具体命令如下:

sudo add-apt-repository ppa:webupd8team/java
sudo apt-get update
sudo apt-get install oracle-java8-installer

安装java8会有图形化向导,按照向导指引安装就好了。

安装了java8以后执行以下命令安装bazel:

echo "deb [arch=amd64] http://storage.googleapis.com/bazel-apt stable jdk1.8" | sudo tee /etc/apt/sources.list.d/bazel.list
curl https://storage.googleapis.com/bazel-apt/doc/apt-key.pub.gpg | sudo apt-key add -
sudo apt-get update && sudo apt-get install bazel
sudo apt-get upgrade bazel

经过以上准备,相关的软件都安装配置好了,可以开始编译TensorFlow了。

首先把TensorFlow项目克隆到本地:

git clone --recurse-submodules https://github.com/tensorflow/tensorflow

接着进入tensorflow目录并运行configure命令为编译做准备,注意要加上sudo,否则会出现权限的问题。

cd tensorflow
 
sudo ./configure

执行./configure会出现文本配置界面,在选择不同选项时几乎可以都使用缺省值,当然在选择是否有GPU支持时要选择Y,然后就是要注意“Cuda compute capability”一项,如果你选择的实例类型是G2,这里需要输入3.0,因为G2实例使用的GPU是K520,只能支持Cuda compute capability 3.0。如果你选择的机型是P2,因为使用的GPU是K80,可以使用缺省的Cuda compute capability,本例使用的TensorFlow版本缺省的是3.5。

具体选项请参考以下截图:

如果你为G2实例类型选择了3.5或者更高版本,运行TensorFlow时会出现以下错误信息:

Ignoring visible gpu device (device: 0, name: GRID K520, pci bus id: 0000:00:03.0) with Cuda compute capability 3.0. The minimum required Cuda capability is 3.5.

配置好以后通过bazel命令编译,注意使用sudo命令:

sudo bazel build -c opt --config=cuda //tensorflow/tools/pip_package:build_pip_package

编译需要一段时间,编译完成后使用buildpippackage命令为python构建pip包,以下命令中的/tmp/tensorflow_pkg目录是pip包输出的临时文件,可以根据你自己的环境进行调整:

sudo bazel-bin/tensorflow/tools/pip_package/build_pip_package /tmp/tensorflow_pkg

构建了pip包以后使用pip命令进行安装,注意以下命令中的/tmp/tensorflow_pkg/tensorflow-0.11.0-cp27-cp27mu-linux_x86_64.whl是上一个命令生成的whl包文件路径,其中包括了具体的版本号,在你自己编译的时候需要留意这个版本号是否有变化:

sudo pip install /tmp/tensorflow_pkg/tensorflow-0.11.0-cp27-cp27mu-linux_x86_64.whl

以上命令执行完以后就成功安装了TensorFlow了,可以开始测试了。

测试与分布式部署

神经网络的测试经常使用mnist数据集进行,TensorFlow中自带了mnist的样例,安装好TensorFlow以后通过以下命令可以进行测试:

python tensorflow/tensorflow/models/image/mnist/convolutional.py

一切正常的话,很快convolutional.py命令就可以执行完成,结束后会输出准确率。

测试结束后,进一步需要做的事情就是分布式部署。如开篇讲到的,在AWS上可以使用镜像(AMI)简化新实例启动过程中的安装工作。

基本做法就是选中以上步骤配置的EC2实例,然后选择“”操作执行镜像创建的操作。当我们需要启动第二台GPU实例时,选择刚创建的镜像进行启动就好了。

启动第二台实例时注意选择第一台实例使用的子网,减少其它因素的干扰。如果对网络性能有特别要求,还可以将两台实例放在同一个“置放组”中,提高两台GPU实例之间的网络通讯速度。

另外,如下文所示,TensorFlow不同进程启动的时候会使用不同的端口,本例使用的是2222和2223端口。为了让两个GPU实例可以通过这些端口通讯,需要设置两个GPU实例所使用的安全组,开放2222,2223,2224的入站访问。

在tensorflow的源代码目录中有个分布式样例文件叫:./tensorflow/tools/dist_test/python/mnist_replica.py,是已经写好的分布式代码。

该python程序通过–pshosts指定参数服务的IP和端口,通过–workerhosts指定工作进程的IP地址和端口。

本例中使用的两台GPU实例的ip地址分布是172.31.24.117和172.31.30.43,我在172.31.24.117上运行参数服务器进程,命令如下:

python ./tensorflow/tools/dist_test/python/mnist_replica.py --ps_hosts=172.31.24.117:2222 --worker_hosts=172.31.24.117:2223,172.31.30.43:2223 --job_name=ps --task_index=0 

注意这里的参数–job_name=ps用于指定本进程的角色。

接着在172.31.24.117上起第一个worker进程,命令如下:

python ./tensorflow/tools/dist_test/python/mnist_replica.py --ps_hosts=172.31.24.117:2222 --worker_hosts=172.31.24.117:2223,172.31.30.43:2223 --job_name=worker --task_index=0 

这个命令和第一个启动参数服务器的命令几乎一样,差别就是–job_name=worker这个参数。

接着在172.31.30.43上启动第二个worker进程,命令如下

python ./tensorflow/tools/dist_test/python/mnist_replica.py –ps_hosts=172.31.24.117:2222 –worker_hosts=172.31.24.117:2223,172.31.30.43:2223 –job_name=worker –task_index=1
因为都是启动worker进程,所以两个启动worker进程的命令都是用的 –jobname=worker 参数,关键在于第一个worker进程使用taskindex=0 参数,而第二个worker进程使用task_index=1 参数。

执行以上命令后稍等片刻两个worker进程就会开始mnist数据集的训练,从日志输出上可以看出两个worker进程分别负责了部分样本的训练。

结语

以上就是AWS上编译安装TensorFlow并执行分布式训练的过程,完成以上环境搭建以后就可以开始你自己的分别式神经网络训练了。

作者介绍:

邓明轩

亚马逊AWS解决方案架构师,拥有15年IT 领域的工作经验,先后在IBM,RIM,Apple 等企业担任工程师、架构师等职位;目前就职于AWS,担任解决方案架构师一职。喜欢编程,喜欢各种编程语言,尤其喜欢Lisp。喜欢新技术,喜欢各种技术挑战,目前在集中精力学习分布式计算环境下的机器学习算法。