GPU版本TensorFlow安装指南

GPU版本的TensorFlow可以利用NVIDIA GPU强大的计算加速能力,使TensorFlow的运行更为高效,尤其是可以成倍提升模型训练的速度。

在安装GPU版本的TensorFlow前,你需要具有一块不太旧的NVIDIA显卡,以及正确安装NVIDIA显卡驱动程序、CUDA Toolkit和cnDNN。

GPU硬件的准备

TensorFlow对NVIDIA显卡的支持较为完备。对于NVIDIA显卡,要求其CUDA Compute Capability须不低于3.0,可以到 NVIDIA的官方网站 查询自己所用显卡的CUDA Compute Capability。目前,AMD的显卡也开始对TensorFlow提供支持,可访问 这篇博客文章 查看详情。

NVIDIA驱动程序的安装

Windows

Windows环境中,如果系统具有NVIDIA显卡,则往往已经自动安装了NVIDIA显卡驱动程序。如未安装,直接访问 NVIDIA官方网站 下载并安装对应型号的最新公版驱动程序即可。

Linux

在服务器版Linux系统下,同样访问 NVIDIA官方网站 下载驱动程序(为 .run 文件),并使用 sudo bash DRIVER_FILE_NAME.run 命令安装驱动即可。在安装之前,可能需要使用 sudo apt-get install build-essential 安装合适的编译环境。

在具有图形界面的桌面版Linux系统上,NVIDIA显卡驱动程序需要一些额外的配置,否则会出现无法登录等各种错误。如果需要在Linux下手动安装NVIDIA驱动,注意在安装前进行以下步骤(以Ubuntu为例):

  • 禁用系统自带的开源显卡驱动Nouveau(在 /etc/modprobe.d/blacklist.conf 文件中添加一行 blacklist nouveau ,使用 sudo update-initramfs -u 更新内核,并重启)

  • 禁用主板的Secure Boot功能

  • 停用桌面环境(如 sudo service lightdm stop

  • 删除原有NVIDIA驱动程序(如 sudo apt-get purge nvidia*

小技巧

对于桌面版Ubuntu系统,有一个很简易的NVIDIA驱动安装方法:在系统设置(System Setting)里面选软件与更新(Software & Updates),然后点选Additional Drivers里面的“Using NVIDIA binary driver”选项并点选右下角的“Apply Changes”即可,系统即会自动安装NVIDIA驱动,但是通过这种安装方式安装的NVIDIA驱动往往版本较旧。

NVIDIA驱动程序安装完成后,可在命令行下使用 nvidia-smi 命令检查是否安装成功,若成功则会打印出当前系统安装的NVIDIA驱动信息,形式如下:

  1. $ nvidia-smi
  2. Mon Jun 10 23:19:54 2019
  3. +-----------------------------------------------------------------------------+
  4. | NVIDIA-SMI 419.35 Driver Version: 419.35 CUDA Version: 10.1 |
  5. |-------------------------------+----------------------+----------------------+
  6. | GPU Name TCC/WDDM | Bus-Id Disp.A | Volatile Uncorr. ECC |
  7. | Fan Temp Perf Pwr:Usage/Cap| Memory-Usage | GPU-Util Compute M. |
  8. |===============================+======================+======================|
  9. | 0 GeForce GTX 106... WDDM | 00000000:01:00.0 On | N/A |
  10. | 27% 51C P8 13W / 180W | 1516MiB / 6144MiB | 0% Default |
  11. +-------------------------------+----------------------+----------------------+
  12.  
  13. +-----------------------------------------------------------------------------+
  14. | Processes: GPU Memory |
  15. | GPU PID Type Process name Usage |
  16. |=============================================================================|
  17. | 0 572 C+G Insufficient Permissions N/A |
  18. +-----------------------------------------------------------------------------+

提示

命令 nvidia-smi 可以查看机器上现有的GPU及使用情况。(在Windows下,将 C:\Program Files\NVIDIA Corporation\NVSMI 加入Path环境变量中即可,或Windows 10下可使用任务管理器的“性能”标签查看显卡信息)

更详细的GPU环境配置指导可以参考 这篇文章这篇中文博客

CUDA Toolkit和cnDNN的安装

在Anaconda环境下,推荐使用

  1. conda install cudatoolkit=X.X
  2. conda install cudnn=X.X.X

安装CUDA Toolkit和cnDNN,其中X.X和X.X.X分别为需要安装的CUDA Toolkit和cuDNN版本号,必须严格按照TensorFlow官方网站所说明的版本安装。在安装前,可使用 conda search cudatoolkitconda search cudnn 搜索conda源中可用的版本号。

当然,也可以按照 TensorFlow官方网站上的说明 手动下载CUDA Toolkit和cuDNN并安装,不过过程会稍繁琐。

使用conda包管理器安装GPU版本的TensorFlow时,会自动安装对应版本的CUDA Toolkit和cuDNN。conda源的更新较慢,如果对版本不太介意,推荐直接使用 conda install tensorflow-gpu 进行安装。