跳到主要内容

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 安装视频演示​