跳到主要内容

TensorFlow 安装 - GPU

与 CPU 版本相比,因为可以利用 GPU 进行更加复杂的运算,因此 GPU 版本的 TensorFlow 要比 CPU 版本的 TensorFlow 强大不少。相应的,GPU 版本的 TensorFlow 安装步骤也更加复杂。那么我们这一节就来看一下 TensorFlow-GPU 应该如何安装。

首先我们要明确的是,如果想要使用 TensorFlow 的 GPU 版本,那么首先拥有 TensorFlow 所支持的 GPU 硬件与相应的软件支持。目前来说 TensorFlow 在官方文档之中给出的要求为:

硬件要求:CUDA® 计算能力为 3.5 或更高的 NVIDIA® GPU 卡。

软件要求

  • NVIDIA® GPU 驱动程序:CUDA 10.1 需要 418.x 或更高版本;
  • CUDA® 工具包:TensorFlow 支持 CUDA 10.1(TensorFlow 2.1.0 及更高版本);
  • CUDA 工具包附带的 CUPTI;
  • cuDNN SDK(7.6 及更高版本);
  • TensorRT 6.0,可缩短用某些模型进行推断的延迟时间并提高吞吐量(可选)。

一般而言,我们 Nvidia® GTX 750 及其以上的显卡都可以达到计算能力3.5 (GTX 760 与 GTX 770除外)。对于软件需求我们会在下面的安装教程之中帮助大家安装与升级。

所以如果大家的 GPU显卡不符合要求,那么就需要更新自己的硬件条件或者采用 CPU 版本的TensorFlow。那么这节课我们就带大家学习以下如何在硬件条件符合的情况之下,在 Windows 与 Linux 上安装 TensorFlow。

1. 在 Windows 上安装 TensorFlow-GPU

在 Windows 上安装 TensorFlow-GPU 可以分为几个步骤,具体来说如下所示:

  • 安装 TensorFlow-GPU;
  • 保证自己的显卡驱动正常;
  • 安装符合自己版本的 CUDA;
  • 安装符合自己版本的 cuDNN 工具包;
  • 上述过程中的各种配置。

既然明确了思路,那么我们就开始在 Windows 上安装 TensorFlow-GPU。

1.1 安装 TensorFlow-GPU

我们在这里依然是使用 pip 进行安装,在进行安装之前需要查看 Python 与 pip 的版本是否符合要求。具体说来,我们一般要求:

  • Python 3.5 及其以上版本;
  • pip 19.0 及其以上版本。

如果版本要求不符合,那么我们则需要进行相应的升级操作。

在确保环境符合条件的情况之下我们可以使用命令安装:

pip install tensorflow-gpu

如果自己的 GPU 驱动版本较低,我们可以指定安装之前的 TensorFlow 版本,如 2.0 版本:

pip install tensorflow-gpu==2.0

安装完成后我们可以进行后面的操作。

1.2 确保自己的显卡驱动正常运行

这一步主要是确保自己的显卡能够正常使用,如果自己的 GPU 驱动过低,那么就需要根据自己的驱动版本进行相应的更新。

目前来说只要将 Nvidia 显卡驱动更新至 410.x 版本及其以上就足够使用了。

1.3 安装符合自己版本的 CUDA 工具包

这一个步骤的版本因人而异,因为每个人的显卡版本不一样,并且每个人的显卡驱动也不尽相同,所以需要根据自己的实际情况而来。故而这一步的操作较为复杂。

首先要查看自己显卡驱动的 CUDA Driver 版本。 打开 Nvidia 控制面板,并点击查看信息,再切换到“组件”选项卡,在此界面可以查看到自己的 CUDA Driver 版本,如下图所示:

图片描述

在上图的示例之中,CUDA Driver 的版本为 11.0 。

去 Nvidia 开发者官方网站下载相应的 CUDA 工具包。 我们可以访问 Nvidia 开发者官网 CUDA 下载,访问的结果如下图所示:

图片描述

然后在这个页面之中选择符合自己版本的 CUDA 下载。 因为上文之中的CUDA Driver 版本为11.0版本,故而我们需要下载第一个 CUDA 工具包。

安装 CUDA 工具包,该步骤较为简单,我们只需要双击打开,并根据指示一步步安装即可。

图片描述

安装完成之后根据提示,添加系统变量:

图片描述

并且在 Path 变量之中添加刚刚新增的变量:%CUDA10% 。

Tips:该处的名字可以自己取,也可以将路径直接添加到PATH之中。

1.4 安装 cuDNN

在这一步,我们要安装 cuDNN 这一个扩展工具,该工具能够很大程度地提升 TensorFlow 的训练速度。具体来说,我们可以按照下面几个步骤进行安装:

登录 cuDNN官方下载网站,并选择自己的版本下载,如下图所示:

我们可以根据自己的 CUDA 版本来选择自己要下载cuDNN版本,比如上文的 CUDA 版本为11.0,那么我们便可以下载下图中的第一个cuDNN版本。

图片描述

下载完成后进行解压,并将解压完成的三个文件夹复制到 CUDA 的安装目录即可。

1.5 测试

安装完成之后,我们可以使用如下代码进行测试:

import tensorflow as tf
print("Version: " + tf.__version__)

若是输出结果为:

Version: 2.2.0

则说明安装成功。

若是没有输出版本并且报错,那么可能是自己的 CUDA 与 cuDNN 的版本与当前的 TensorFlow 版本不匹配,需要重新安装。一般而言,如果要安装 TensorFlow-GPU 2.x , 我们需要安装 7.4 及其以上版本的 cuDNN 以及 10.0 及其版本以上的 CUDA 。

2. 在 Linux 上安装 TensorFlow-GPU

在这一部分,我们会在命令行的环境下在 Linux 系统之中安装 TensorFlow 的 GPU 版本。具体步骤分为以下几点:

2.1 查看 Nvidia 驱动是否安装

我们可以采用以下命令来查看 Nvidia 的驱动是否安装:

nvidia-smi

如果该命令成功地输出了 Nivida 显卡地相关信息,则表示驱动成功安装;否则的话需要安装或者更新 Nvidia 显卡驱动。

如果需要更新驱动,我们可以通过如下命令来卸载原来的驱动:

sudo apt-get remove --purge nvidia*

然后我们下载好自己需要更新的版本,然后便可以安装新的驱动:

sudo service lightdm stop
sudo ./NVIDIA-Linux-****.run
sudo service lightdm start

2. 安裝 CUDA 工具

首先我们需要像 Windows 安装一样,找到自己的 CUDA 版本。

cat /usr/local/cuda/version.txt

Tips:不同用户的驱动的安装位置不一样,这只是默认的安装路径。

然后根据自己的 CUDA Driver 的版本,找到适合自己 CUDA Driver 版本的 CUDA 工具包,然后我们便可以运行以下命令来进行安装。

wget http://developer.download.nvidia.com/compute/cuda/11.0.2/local_installers/cuda_11.0.2_450.51.05_linux.run

sudo sh cuda_11.0.2_450.51.05_linux.run

这只是安装的一种方式,其实官方网站给我们提供了非常完备的安装方案,每当我们选择一个版本与安装方式之后,官方便会给我们详细的安装指引,从而指引我们完成安装。

图片描述

配置环境变量:

export PATH=/usr/local/cuda-11.0/bin:$PATH
export LD_LIBRARY_PATH=/usr/local/cuda-11.0/lib64:$LD\_LIBRARY\_PATH

当然,这里面的路径是你自己的 CUDA 的安装路径

2.3 安装 cuDNN

与 Windows 的安装一样,我们都需要使得 CUDA 与 CuDNN 的版本相对应,因此我们要首先要确定自己的要下载的 cuDNN 版本并且下载到自己的机器上(这里以 tar 文件为例)。

运行命令,解压 cuDNN:

tar -xzvf cudnn-x.x-linux-x64-v8.x.x.x.tgz

将文件拷贝到 CUDA 的安装目录

sudo cp cuda/include/cudnn*.h /usr/local/cuda/include
sudo cp cuda/lib64/libcudnn* /usr/local/cuda/lib64

如果之后运行失败的话可能是权限问题,需要提升以下权限。

我们可以通过以下命令来提升一下目录的权限。

chmod -R 777 /usr/local/cuda

2.4 安装 TensorFlow-GPU

与之前一样,我们可以运行以下命令来使用pip安装 tensorflow-gpu:

pip install tensorflow-gpu

2.5 测试是否成功

安装完成之后,我们可以运行代码:

import tensorflow as tf
print("Version: " + tf.__version__)

若是输出结果为:

Version: 2.2.0

则说明安装成功。

3. 附加说明

有些同学可能会问“为什么没有安装 CUPTI ?这个不是必须的吗” ?

没错,这个 CUPTI 工具是必须要安装的,只是我们在安装CUDA的时候已经附带安装,因此不需要重新安装。

如果大家遇到下载速度较慢的问题便可以尝试更换国内的下载源:

pip install tensorflow-gpu -i https://pypi.tuna.tsinghua.edu.cn/simple

4. TensorFlow-GPU 安装视频演示