在 Windows 平台安装 NVIDIA GPU 加速的 TensorFlow

TensorFlow 支持使用 NVIDIA GPU 进行加速,但相关文档比较分散,较难操作。本文整理了在 Windows 平台上安装和配置 TensorFlow 并使用 NVIDIA GPU 进行加速的方法。使用环境为 Windows 10 x64 + NVIDIA GeForce GTX 970。

0. 显卡要求

要使用 NVIDIA GPU 加速 TensorFlow,首先要有支持 CUDA Compute Capability 3.0 以上的 NVIDIA GPU 设备。CUDA 是 NVIDIA 推出的一套并行计算平台和 API,NVIDIA 各产品的支持情况可以在这里找到。以 GeForce GTX 970 显卡为例,在 CUDA-Enabled GeForce Products 下面可以找到其 Compute Capability 为 5.2,如图 1,满足 TensorFlow 的要求。

图 1

图1

1. 安装 Python 3

在 Windows 平台上,TensorFlow 仅支持 Python 3.5.x,可以直接从官方下载。由于我之前一直使用 Anaconda,下面再介绍一下 Anaconda 的安装方法。

按照官方的定义,Anaconda 是一个基于 Python 的数据科学平台,它集成了一系列常用的科学计算包,同时提供了一个包管理工具 conda,在很大程度上解决了 Python 各种包多如牛毛难以选择和管理的问题。

安装 Anaconda 非常简单,在这里 选择对应的系统和 Python 版本,下载并安装即可,安装过程中还可以选择配置环境变量,非常方便。Anaconda 提供了 Python 2 和 Python 3 两个版本,本文编写时,Anaconda 的 Python 3 版本使用的是 Python 3.6,而 TensorFlow 仅支持到 Python 3.5.x(后续会支持 3.6),后面会介绍建立 Python 3.5 环境的方法。

2. 安装 CUDA Toolkit 和 cuDNN

2.1. 安装 NVIDIA CUDA Toolkit

NVIDIA CUDA Toolkit 是一个面向 C/C++ 开发者的用于开发 GPU 加速应用程序的环境,可以在这里下载。选择平台后下载到一个安装程序,直接安装即可,目前最新版本为 CUDA Toolkit 8.0

2.2. 安装 NVIDIA cuDNN

NVIDIA cuDNN (CUDA Deep Neural Network) 是一个用于深度神经网络的 GPU 加速库,它的安装比较麻烦,在这里,注册、登录并同意 License 后,可以看到一系列 cuDNN 的版本,如图 2:

图 2

图 2

这里需要根据上一步安装的 CUDA Toolkit 版本和 TensorFlow 所支持的版本来选择。上一步安装的 CUDA Toolkit 版本为 8.0,从 TensorFlow 的安装文档得知目前支持的 cuDNN 版本为 5.1,综合二者,需要选择 Download cuDNN v5.1 (Jan 20, 2017), for CUDA 8.0

下载后得到一个 zip 文件,解压后,得到一个名为 cuda 的文件夹,把 cuda 和其中的 cuda\bin 两个路径加入 PATH 环境变量即可。

3. 安装 TensorFlow

终于到正题了。本文主要介绍把 TensorFlow 安装到 Anaconda Environment 的方法。Anaconda Environment 可以看做是一个 Python 环境的沙盒,不同的 Environment 可以包含不同版本的 Python 和各种不同的包,通过 Environment 可以独立地为各个的项目配置环境和依赖,互不影响。

3.1. 创建 Envrionment

首先创建一个用于 TensorFlow 的 Environment。执行:

conda create -n tensorflow python=3.5 anaconda nb_conda

创建一个名为 tensorflow 的环境。前面提到,目前 TensorFlow 在 Windows 上仅支持 Python 3.5.x,所以这里使用 python=3.5 指定了 Python 版本。anacondanb_conda 为要预装的两个包。anaconda 包内集成了各种常用的科学计算包,nb_conda 则是用于切换 Jupyter Notebook 的 Kernel 环境,后面会做说明。

由于包含了 anaconda 包,安装时间较长。也可以使用诸如

conda create -n tensorflow python=3.5 numpy scipy matplotlib nb_conda

的命令,仅包含自己常用的包,之后如有其它需要,再单独安装。

3.2. 安装 TensorFlow

启动前面建立的 tensorflow 环境:

activate tensorflow

再安装 GPU 版本的 TensorFlow:

pip install --ignore-installed --upgrade https://storage.googleapis.com/tensorflow/windows/gpu/tensorflow_gpu-1.0.1-cp35-cp35m-win_amd64.whl

这里建议不要直接使用上面命令中的版本,而是去官方文档查看最新版本,或者使用 GitHub 页面 上的 Nightly Build 版本。我在进行这一步的时候,使用官方文档的版本可以正常安装,但运行时会有诸如“OpKernel (‘op: “BestSplits” device_type: “CPU”‘) for unknown op: BestSplits”的报错,可以查到相关 Issue,使用了 Nightly Build 版本才正常。

3.3. Hello, TensorFlow

惯例使用 Hello 系列验证安装。还是在 tensorflow 环境,进入 Python,输入:

import tensorflow as tf


hello = tf.constant('Hello, TensorFlow!')
sess = tf.Session()
print(sess.run(hello))

如果运行失败,去搜下 Issue 把,说不定已经在 Nightly Build 修复了。

4. 使用 Jupiter Notebook

最后再介绍一下在 Anaconda Environment 中使用 Jupiter Notebook 的方法。

前面把 TensorFlow 安装进了名为 tensorflow 的 Anaconda Environment 中,但 Jupiter Notebook 本身不支持 Anaconda Environment 的切换,也就无法使用 tensorflow 环境中的 TensorFlow。在 3.1. 中安装的 nb_conda 包就是为了解决这个问题,安装了 nb_conda 之后,在 Jupiter Notebook 里会出现一个 Conda 标签,可以查看所有的 Anaconda Environment,如图 3:

图 3

图 3

也可以在 Kernel / Change kernel 下切换当前文档所使用的环境,如图 4:

图 4

图 4