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

本文转载自 AWS 技术博客。

原文链接：

https://amazonaws-china.com/cn/blogs/china/tensorflow-on-aws/