TensorFlow on Android（10）: 在GPU上进行训练

一般来说我们会在笔记本或者PC上面编写模型和训练代码，准备一些数据，配置训练之后会在笔记本或者PC上面做一个简单验证，如果这些代码数据都OK的话， 然后真正的训练放在计算力更强的的计算机上面执行，一般来说至少有一块或者多块CPU， 有相当的显存和内存，接下来我们就这样实验一下。

选择一个支持TensorFlow GPU的计算机

当务之急是找到一块可以用于TensorFlow的显卡，TensorFlow只支持在NVIDIA的部分高端显卡上面进行GPU加速， 在NVIDIA开发者中心可以找到支持的显卡列表，确保你的显卡在这个列表之内。

一般来说有两个选择， 第一你可以自己购买配置一个专门用于机器学习/深度学习的服务器，这样初期的投入比较高。 第二可以租用云服务器，目前各大云服务商都有专门的GPU计算实例，可以按照小时计费，这样对于学习来说成本还是可以接受的。

其实的我的Macbook Pro上面有一块N卡， 但是从TensorFlow 1.2开始， 官方已经不再支持Mac上的GPU计算了。 虽然可以通过一些hack使Tensorflow的Mac版本继续支持GPU，但是笔记本上的显卡计算能力还是比较弱，我也不想训练到一半把这块显卡烧了，所以我选择从云服务商那里租用一台GPU计算实例。

由于暂时没有人付我广告费，所以在这里就不透露我选择的是哪家云服务商的产品了 ^_, 总之我租用了一台装有NVIDIA K80的，安装Ubuntu 16.04.1 LTS的主机，综合的费用的话大概不到5块人民币一小时。接下来在这台主机上面继续训练。

安装CUDA

CUDA是一种由NVIDIA推出的通用并行计算架构， 只能应用于NVIDIA的系列显卡， 目前主流的深度学习框架包括TensorFlow都是使用CUDA来进行GPU加速的（可以把CUDA当做是一种驱动程序， TensorFlow通过这个驱动程序来使用显卡），所以我们必须先安装CUDA。

首先我们先检查一下显卡信息，确保没有租错主机^_, 在终端上运行：

lspci | grep -i nvidia

你要确保列出的显卡在NVIDIA支持CUDA的显卡列表里面。

在我租用的主机上，显示如下：

显卡没有问题，接下安装CUDA(本系列文章使用CUDA 8)：

在NVIDIA开发者中心下载相应的deb包。

下载完成以后运行：

sudo dpkg -i cuda-repo-ubuntu1604-8-0-local-ga2_8.0.61-1_amd64.deb
sudo apt-get update
sudo apt-get install cuda

如果不出意外，安装完成之后在终端运行：

nvcc --version

会输出CUDA的版本信息

安装cuDNN

我们还需要安装一个叫做cuDNN的深度卷积神经网络的加速库：
在NVIDIA开发者中心下载安装包，这需要你注册一个账号。

注意：选择cuDNN6的版本， 这样才和CUDA8兼容

下载完成后运行：

sudo tar -zxf cudnn-8.0-linux-x64-v6.0.tgz 
cd cuda
sudo cp lib64/* /usr/local/cuda/lib64/
sudo cp include/cudnn.h /usr/local/cuda/include/

然后编辑~/.bashrc, 在最后面加入：

export CUDA_HOME=/usr/local/cuda
export PATH=${CUDA_HOME}/bin/:$PATH
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$CUDA_HOME/lib64:$CUDA_HOME/extras/CUPTI/lib64

然后source ~/.bashrc， 安装完成。

安装TensorFlow GPU版

为了在GPU上进行训练，我们还要安装TensorFlow的GPU版本（之前我们在笔记本上面安装的而是CPU版）：

sudo pip install tensorflow-gpu==1.3.0

安装完成以后，运行Python：

python
>>> from tensorflow.python.client import device_lib
>>> print device_lib.list_local_devices()

如果输出信息里面有/gpu:0的信息， 那么说明整个安装就成功完成了。

继续训练

前面化了一点时间来配置GPU的环境，现在是时候继续我们的训练了。
当然我们还是需要在这台机器上面根据上一课时的内容完成Object Detection API的安装和配置；下载Pre-trained 模型， 然后把我们本地的训练目录打包上传， 然后根据具体的路径修改pipeline.config中的相应项； 然后运行训练脚本，这次的输出是这样的：

每一步的时间被缩短到1秒左右， 那么执行完200000步大概需要2天多一点时间， 对于实际应用，其实是可以接受的。

我在运行到20000多步的时候终止了学习（大概6个多小时）, 然后将训练目录拷贝到了本地：

在model/train目录下就是这20000多步的学习输出， 如果你在这个目录上继续运行训练脚本的话，它会从之前的学习进度开始，继续学习。

我将训练好的模型导入到《TensorFlow on Android：物体识别》中的Android App中体验了一下效果：

效果还不错！

---

一个训练的流程就跑完了，我们配置好了深度学习的软硬件环境，接下来我们开始准备数据，训练自己的模型吧。