在AWS GPU主机上用Docker容器运行TensorFlow

学习TensorFlow,第一步当然要把TensorFlow的程序跑起来。在什么样的环境上跑呢?深度学习的训练包含大量的浮点运算,业内广泛使用GPU来进行加速,加速比可以达到10到100倍。我们作为专业人士,当然也要用GPU来加速训练。本文就教大家怎么用最经济快速的方式使用GPU来运行TensorFlow。

技术栈

深度学习使用的GPU都比较贵,业内典型的GPU,Titan X K40,售价大概20000RMB。一次性投入过大,不适合短期使用。目前很多云提供商都提供了带GPU的云主机,支持按使用量付费,比较适合短期的使用。我们将选择带GPU的云主机作为我们的环境。同时,Docker是我们一直非常喜欢的运行程序的方式,TensorFlow支持跑在Docker容器里,并且能够使用GPU的运算能力。但要做到这一点,需要在我们的运行环境中增加一些层级和组件。为了大家有个整体的认识,我们先介绍以下我们的技术栈。

1. 云主机+GPU

国内提供带GPU的云主机的厂商有阿里云和UCloud。不幸的是目前这两家都缺货,阿里云支持在控制台页面直接申请,但是提示缺货,正在备货。这个状态已经持续超过一个月了。UCloud需要申请,据客服说申请者要等至少半个月。国外的亚马逊AWS也提供带GPU的主机,现货充足。我们就以这个环境为我们的运行环境。

2. Docker

Docker的好处在于通过它运行程序非常方便,因为所有的依赖都包含在一个镜像当中了,一条命令就可以把程序跑起来,也可以将自己的程序打包成镜像到别的机器上很方便地运行。对于Docker不熟悉的同学,可以把Docker容器理解成轻量级的虚拟机,或者彼此隔离的进程/进程树。关于Docker的进一步介绍可以看这里

3. NVIDIA驱动及Docker Plugin

NVIDIA的GPU在深度学习领域占有近80%的市场,亚马逊的云主机所配的也是NVIDIA的GPU。要使用NVIDIA GPU,我们需要安装NVIDIA驱动,同时,因为我们是通过Docker容器的方式来运行TensorFlow,我们还需要一个额外的东西NVIDIA Docker Plugin。NVIDIA Docker Plugin把主机上的GPU能力暴露给容器,使得容器也能使用GPU。

4. NVIDIA cuDNN库

为了更方便GPU的用户开发深度学习程序,更有效地使用GPU的计算能力,NVIDIA提供了cuda写的DNN(Deep Neural Network)库。Tenforflow的GPU版也使用了这个库,所以我们还需要在容器内安装这个库。

5. 清单

综合上面的内容,我们的运行技术栈如下(按从底层往上层顺序排列):

  • AWS g2.8xlarge型号主机
  • Amazon Linux
  • NVIDIA driver
  • Docker
  • NVIDIA Docker plugin
  • cuDNN库
  • Tensorflow GPU版

如下图所示:

步骤

明确了技术栈之后,我们要做的是把整个技术栈搭建起来并在上面跑TensorFlow神经网络。具体的执行步骤如下。

1. 申请云主机

到亚马逊AWS网站申请带GPU的机器,高配型号g2.8xlarge,低配型号g2.2xlarge。我们使用的是g2.8xlarge。这里的步骤就不详细介绍了。网上有很多的AWS的教程。需要注意的是,亚马逊的云主机支持各种Linux版本,我们推荐使用Amazon Linux。RedHat和CentOS在使用中遇到包括系统启动时间过长,或安装NVIDIA驱动的时候提示内核不兼容等问题。

2. 安装NVIDIA driver

到NVIDIA网站 http://www.nvidia.com/object/unix.html 下载./NVIDIA-Linux-x86_64-367.35.run,改属性,然后直接运行,会出现文本图形界面,按照界面提示进行安装。

sudo yum install gcc  
sudo yum install kernel-devel-$(uname -r)  
sudo chmod +x ./NVIDIA-Linux-x86_64-367.44.run  
sudo ./NVIDIA-Linux-x86_64-367.44.run #注意要选32-bit compatible和X  

(参考:https://github.com/NVIDIA/nvidia-docker/wiki/Installation)

3. 安装Docker

sudo yum -y update  
sudo yum -y install docker  

注意确保Docker容器能正常运行,如果有问题,参考Trouble Shooting

4. 安装NVIDIA Docker以及NVIDIA Docker plugin

参考这里的步骤,安装NVIDIA Docker:Install instruction for other distributions(https://github.com/NVIDIA/nvidia-docker/tree/v1.0.0-rc.2#other-distributions))

# Install nvidia-docker and nvidia-docker-plugin
wget -P /tmp https://github.com/NVIDIA/nvidia-docker/releases/download/v1.0.0-rc.2/nvidia-docker_1.0.0.rc.2_amd64.tar.xz  
sudo tar --strip-components=1 -C /usr/bin -xvf /tmp/nvidia-docker*.tar.xz && rm /tmp/nvidia-docker*.tar.xz

# Run nvidia-docker-plugin
sudo -b nohup nvidia-docker-plugin > /tmp/nvidia-docker.log

# Test nvidia-smi
sudo nvidia-docker run --rm nvidia/cuda nvidia-smi  

这一步很重要,我们需要理解这里做了什么,并确保成功运行。这一组命令做了三件事情:
1. 下载和安装NVIDIA docker和NVIDIA docker plugin;
2. 运行NVIDIA docker plugin;
3. 通过NVIDIA docker运行nvidia/cuda容器,跑nvidia-smi命令。

NVIDIA docker可以理解为NVIDIA在docker外面包围的一个“壳”,之后都要用NVIDIA docker来启动docker容器,这样容器才能访问到GPU。

倒数第二条命令启动NVIDIA docker plugin。建议打开/tmp/nvidia-docker.log查看nvidia-docker-plugin启动的结果,应该不显示错误。如果有错误,参考文末的常见问题和解决方案进行修复。常见的一个错误是"Could not load UVM kernel module",一般来说,是NVIDIA driver安装不正确,需要重新安装(注意选择32bit compatible和X配置)。

最后一条命令通过nvidia/cuda容器运行nvidia-smi这个命令(详情参考nvidia-smi命令 )显示GPU设备的状态,成功运行的话会显示如下输出:

这一步成功,就意味着容器已经能够访问GPU了。容器nvidia/cuda能够访问,tensorflow的容器就也应该能够访问了。 参考:https://github.com/NVIDIA/nvidia-docker

Nvidia docker背后原理:https://github.com/NVIDIA/nvidia-docker/wiki/Why%20NVIDIA%20Docker

5. 用NVIDIA docker启动TensorFlow容器

现在我们终于可以启动TensorFlow容器了,注意我们要启动的是支持GPU加速的版本,所以要带上:latest-gpu的tag:

sudo nvidia-docker run -it -p 8888:8888 --name tensorflow gcr.io/tensorflow/tensorflow:latest-gpu  

启动完成后通过docker ps验证启动是否成功:

sudo nvidia-docker ps  

6.安装cuDNN

a. 从NVIDIA网站(https://developer.nvidia.com/cudnn)下载文件,并拷贝到TensorFlow容器的/notebooks目录下。

sudo nvidia-docker cp ./cudnn-7.0-linux-x64-v4.0-prod.tgz tensorflow:/notebooks  

b. 然后进入容器

sudo nvidia-docker exec -it tensorflow bash  

c. 解压和拷贝文件

tar -zxvf ./cudnn-7.0-linux-x64-v4.0-prod.tgz  
mkdir /usr/local/cuda/include  
sudo cp cuda/include/cudnn.h /usr/local/cuda/include  
sudo cp cuda/lib64/libcudnn* /usr/local/cuda/lib64  
sudo chmod a+r /usr/local/cuda/include/cudnn.h /usr/local/cuda/lib64/libcudnn*  

拷贝完成后,应该可以在/usr/loca/cuda/lib64目录下看到cuDNN对应的库文件:

d. 配置包路径

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

7. 运行TensorFlow的CNN模型

一切准备就绪,我们可以运行TensorFlow的CNN模型了。还是那个经典的MNIST手写字符数据集,我们跑TensorFlow自带的CNN网络。

a. 进入容器

sudo nvidia-docker exec -it tensorflow bash  

b. 然后在容器内运行

python -c 'import os; import inspect; import tensorflow; print(os.path.dirname(inspect.getfile(tensorflow)))'  
python -m tensorflow.models.image.mnist.convolutional  

运行结果如下 可以看到,我们在MNIST数据集上成功地跑完了TensorFlow自带的CNN卷积神经网络的训练过程。

总结

现在我们来看以下,GPU的加速比究竟如何。前面说过,GPU能提供10到100倍的加速比。在我们这次实际的例子当中,我们使用的是亚马逊的G2.8xlarge机器,配有4个NVIDIA GRID GPU,每个GPU带有1536个CUDA核心和4G显存(详细的配置参考这篇文章:New G2 Instance Type with 4x More GPU Power)。在MNIST数据集上跑经典卷积网络,普通机器(阿里云主机)上训练耗时6000秒,G2.8xlarge上耗时210秒。加速比29倍,详情如下图。对于深度学习的从业者来说,确实提供了非常大的价值。

常见问题和解决方案

1. 在Amazon Linux上安装docker,启动容器时遇到 'Error starting daemon: Devices cgroup isn't mounted' 错误

这个错误为cgroup在宿主机上没有挂载。我们编辑挂载上去

vi /etc/fstab  
#在结尾添加
none        /sys/fs/cgroup        cgroup        defaults    0    0  

然后执行reboot重启服务器,再尝试启动容器,应该可以成功了。参考: http://www.cnblogs.com/jackluo/p/5441888.html

2. 在执行nvidia-docker-plugin的时候报"Could not load UVM kernel module"

重新安装NVIDIA driver(包括gcc, kernel-devel-),注意在程序提示是否安装32bit compatible以及是否对X进行配置的时候,选择Yes。再执行就好了

3. 在安装NVIDIA driver的时候报找不到gcc,或者kernel source

先运行yum命令安装gcc和kernel source

yum install gcc  
yum install kernel-devel-$(uname -r)  

参考资料

  1. Installing TensorFlow on an AWS EC2 Instance with GPU Support
  2. TensorFlow doc: Installing Cuda GPUs on Linux
  3. NVIDIA Docker:Install instruction for other distributions
  4. NVIDIA Docker背后原理:Why NVIDIA Docker
  5. Running tensorflow on GPU cluster in virtualenv
  6. Add a folder to the Python library path, once for all
  7. AWS实例生命周期
  8. Nvidia GPU + CoreOS + Docker + TensorFlow = A Fast, Flexible, Deep Learning Platform
  9. New G2 Instance Type with 4x More GPU Power
  10. tensorflow的Virtualenv安装方式安装
  11. nvidia-docker | 2016/06/06 22:49:17 Error: Could not load UVM kernel module #106
  12. Could not load nvidia-uvm
  13. [Solved] modprobe can't load nvidia-uvm
  14. nvidia-smi命令
Author image
关于作者 Yijun Wu
Yijun Wu,互联网从业者,信奉用技术手段解决非技术问题。相信人类水平的人工智能在30年之内将会实现。关注深度学习,TensorFlow。