TensorFlow从零开始学
上QQ阅读APP看本书,新人免费读10天
设备和账号都新为新人

2.4 使用GPU 加速

2.4.1 安装配置GPU环境

1.安装GPU版TensorFlow

在2.2节中我们已经介绍了如何安装CPU版TensorFlow,为了使用GPU来加速计算,我们必须安装GPU版TensorFlow。TensorFlow官网对于Docker的安装方法有较为详细的介绍,感兴趣的读者可以参考官网的教程。本节里我们将介绍如何自己手工来安装和配置相关环境。首先使用pip命令:

安装完成后,我们可以查看一下当前可用的GPU:

由于作者的机器上有两块GPU,所以输出两块GPU的编号:“[‘/device:GPU:0’,‘/device:GPU:1’]”。

2.安装显卡驱动

根据你的显卡型号到官网下载对应驱动。图2-23所示的是作者机器Tesla P100的显卡对应的驱动程序。这里一定要注意选择正确的版本,要和你的显卡版本、操作系统版本及想要安装的CUDA版本一一对应(关于TensorFlow与CUDA的版本对应关系,在后面CUDA的安装部分有说明)。

图2-23 NVIDA驱动下载列表

在图2-23所示的界面中单击“搜索”按钮,弹出如图2-24所示的界面,单击“下载”按钮即可。

图2-24 NVIDA驱动下载提示

安装完成之后可以使用“nvidia-smi”命令查看显卡,图2-25所示的是作者机器上的两块显卡的信息。

图2-25 作者机器上的两块显卡的信息

3.安装CUDA

在安装CUDA之前,我们一定要先搞清楚TensorFlow各个版本与CUDA版本的对应关系。在TensorFlow官网有说明,其中Linux系统环境下TensorFlow GPU与CUDA的版本对应关系如图2-26所示。

图2-26 TensorFlow与CUDA的版本对应关系

TensorFlow2.0GPU依赖的NVIDA软件包的官网说明,如图2-27所示。

图2-27 TensorFlow2.0GPU依赖的NVIDA软件包

(1)下载CUDA

首先到NVIDIA官网下载CUDA。作者在撰写本节内容时,CUDA的最新版本是10.1版本,这里提醒读者,一定要按照TensorFlow官网的说明下载10.0版本,否则TensorFlow即使安装好后也是不能正常运行的。后续新的版本可能会支持更高版本的CUDA,读者请根据实际情况下载。

如图2-28所示,选择对应系统环境的CUDA版本。

图2-28 选择对应的CUDA版本

(2)安装CUDA

CUDA下载页面有安装指引,如图2-29所示。

图2-29 CUDA的安装步骤

第一步:执行安装命令。

第二步:添加key。

第三步:依次执行。

注意不要使用“sudo apt-get install cuda”,这样默认安装的是最新版,所以一定要指定版本。安装过程中如果有报“libkmod:ERROR”相关错误,那么安装完成后重启一下即可。

安装完成后,在“/usr/local”目录下会生成“cuda”和“cuda-10.0”两个文件夹,如图2-30所示,我们可以使用命令“cat /usr/local/cuda/version.txt”查看CUDA版本。

图2-30 查看CUDA版本

第四步:设置环境变量。

打开“~/.bashrc”文件,在文件的最后添加如下内容:

在终端执行命令“source~/.bashrc”让环境变量生效。

第五步:验证安装是否成功。

① 进入目录“/usr/local/cuda-10.0/samples/1_Utilities/deviceQuery”中打开终端。

② 终端下执行编译命令“sudo make”。

③ 然后执行命令“./deviceQuery”,可以看到两块GPU的信息。

如图2-31所示,检测到作者的两块显卡,图2-31中所示的是其中一块显卡的信息。到这里CUDA已经安装完成了。

图2-31 显卡的信息

4.安装cuDNN

(1)下载

这里一定要下载与CUDA10.0对应的版本,如图2-32所示。下载cuDNN需要登录NVIDIA账号,没有的话,可以按照提示创建一个账号。

图2-32 cuDNN与CUDA的版本对应关系

选择好cuDNN版本后,单击下载“cuDNN Library for Linux”,如图2-33所示。

图2-33 cuDNN下载列表

(2)安装

第一步:使用“tar”命令解压文件。

第二步:拷贝文件,并修改文件权限。

到这里TensorFlow2.0的GPU版就安装配置完成了。

2.4.2 使用TensorFlow-GPU

如果我们的机器上安装配置好了GPU版的TensorFlow,那么运行的时候TensorFlow会自行去选择可用的GPU,也可以通过“os.environ["CUDA_VISIBLE_DEVICES"]”来选择我们要使用的GPU:

在第4行代码中,我们选择了编号为“0”的GPU,执行完上面的这段代码后使用“nvidia-smi”命令来查看一下GPU的占用情况,如图2-34所示,编号为“0”的GPU正在被占用。我们可以将第4行代码中的“0”改为“1”来使用另一个GPU。

图2-34 查看GPU占用情况

如果我们希望使用多个GPU,例如同时使用“0”“1”两个GPU,可以设置“os.environ ["CUDA_VISIBLE_DEVICES"] = "0,1"”,除此之外还可以使用TensorFlow为“tf.keras”提供的分布式训练策略“tf.distribute.MirroredStrategy”来实现单机环境下的多GPU训练: