如何在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。喜欢新技术，喜欢各种技术挑战，目前在集中精力学习分布式计算环境下的机器学习算法。