Windows 10 上 TensorFlow GPU 2.3.1 配置与安装教程

最近老王换了一台电脑，在训练模型时发现 TensorFlow 用的居然是 CPU 而不是 GPU，后来发现是 CUDA 和 cuDNN 版本的问题。本文记录下 Windows 10 上 TensorFlow GPU 版配置与安装教程，PyCharm 为开发环境。

一、TensorFlow 安装

直接用 pip 安装 TensorFlow：

pip install tensorflow

本文基于 TensorFlow 2.3.1 介绍 GPU 支持和安装步骤。

如果是 1.15 和更早的 TensorFlow 版本，那么 GPU 和 CPU 是分开的，除了 tensorflow 外，还需要安装 tensorflow-gpu：

pip install tensorflow==1.15 # CPU
pip install tensorflow-gpu==1.15 # GPU

二、硬件要求

都开始训练模型了，一般硬件都早已准备好了，这里直接拷贝 TensorFlow 官方对于硬件要求的介绍：

TensorFlow 支持以下带有 GPU 的设备：

1. CUDA® 架构为 3.5、3.7、5.2、6.0、6.1、7.0 或更高的 NVIDIA® GPU 卡。请参阅支持 CUDA® 的 GPU 卡列表。
2. 在配备 NVIDIA® Ampere GPU（CUDA 架构 8.0）或更高版本的系统上，内核已从 PTX 经过了 JIT 编译，因此 TensorFlow 的启动时间可能需要 30 多分钟。通过使用“export CUDA_CACHE_MAXSIZE=2147483648”增加默认 JIT 缓存大小，即可将此系统开销限制为仅在首次启动时发生（有关详细信息，请参阅 JIT 缓存）。
3. 对于 CUDA® 架构不受支持的 GPU，或为了避免从 PTX 进行 JIT 编译，亦或是为了使用不同版本的 NVIDIA® 库，请参阅在 Linux 下从源代码编译指南。
4. 软件包不包含 PTX 代码，但最新支持的 CUDA® 架构除外；因此，如果设置了 CUDA_FORCE_PTX_JIT=1，TensorFlow 将无法在旧版 GPU 上加载。（有关详细信息，请参阅应用兼容性。）

三、软件要求

要在 Windows 上跑 TensorFlow GPU 版，那么必需要有：NVIDIA GPU 驱动程序、CUDA 工具包、CUDA 工具包附带的 CUPTI 和 cuDNN，下面分别介绍每个软件包的安装教程。

注意，在安装 CUDA 和 cuDNN 时，一定要注意选择与 TensorFlow 适配的版本，例如老王这台电脑上的 TensorFlow 是 2.3.1，选择的 CUDA 是 10.1，cuDNN 是 7.6，详细的 tensorflow CUDA cuDNN 版本对应关系介绍可以参考：https://www.tensorflow.org/install/source_windows。

1、NVIDIA GPU 驱动程序安装

直接搜对应的驱动程序就可以了（地址：https://www.nvidia.com/drivers），这个一般装好 GPU 都有，我们在“设备管理器”->“显示设备器”里可以看到自己电脑是否安装了驱动程序以及驱动程序的版本：

2、CUDA 工具包以及 CUPTI 安装

TensorFlow 2.1.0 版本之后开始支持 CUDA 10.1，所以我这直接下的 10.1 版本的 CUDA 工具包：https://developer.nvidia.com/cuda-toolkit-archive，在安装 CUDA 工具包时，CUPTI 应该默认就安装好了，无需额外配置。

3、cuDNN SDK 安装

下载地址：https://developer.nvidia.com/cudnn

cuDNN SDK 下载前需要注册一个 NVDIA 的账号，然后有比较多的问卷需要填写，大致就是问用这个 SDK 干啥，自己的职业是什么，以前是否跑过 ML 等等。注册好后就可以下载了，这里也需要注意版本的问题。

四、环境配置

安装好 CUDA 和 cuDNN 后，需要将它们的路径写入环境变量的 Path 中，根据自己的安装路径，要新增 4 条 Path 数据，老王这里 CUDA 直接是默认的安装路径，cuDNN 则安装在了 D 盘的 tools 文件夹下：

C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v10.1\bin
C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v10.1\extras\CUPTI\lib64
C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v10.1\include
D:\tools\cuda\bin

五、检测 GPU 是否启用

CUDA 和 cuDNN 配置好后，TensorFlow 是会自动调用 GPU 的，不需要额外配置。

这里是 2 个可以进行 TensorFlow GPU 测试的代码，看看 TensorFlow 是在用 CPU 还是在用 GPU：

from tensorflow.python.client import device_lib
import tensorflow as tf
print("Num GPUs Available: ", len(tf.config.experimental.list_physical_devices('GPU')))
print(device_lib.list_local_devices())

如果正确配置了 TensorFlow GPU，那么返回如下：

Num GPUs Available: 1

[name: "/device:CPU:0"
device_type: "CPU"
memory_limit: 268435456
locality {
}
incarnation: 13434866154275638089
, name: "/device:XLA_CPU:0"
device_type: "XLA_CPU"
memory_limit: 17179869184
locality {
}
incarnation: 11984980871348378142
physical_device_desc: "device: XLA_CPU device"
, name: "/device:GPU:0"
device_type: "GPU"
memory_limit: 6507127385
locality {
bus_id: 1
links {
}
}
incarnation: 4195608967092838440
physical_device_desc: "device: 0, name: GeForce RTX 2070 SUPER, pci bus id: 0000:07:00.0, compute capability: 7.5"
, name: "/device:XLA_GPU:0"
device_type: "XLA_GPU"
memory_limit: 17179869184
locality {
}
incarnation: 6938345871348754940
physical_device_desc: "device: XLA_GPU device"
]

当然，我们也可以在跑训练代码时，看任务管理器中的 CPU 和 GPU 使用情况来判断 TensorFlow 是否在用 GPU：