在Windows下安装gpu版的TensorFlow

Posted on Edited on In Views: Word count in article: 487 Reading time ≈ 2 mins.

使用非gpu版本的TensorFlow进行训练实在是太慢了，于是尝试在Windows下安装gpu版本的TensorFlow，写篇日志记录一下。

工作环境

OS：Win10
Python:3.5(目前TensorFlow仍不支持3.6)

安装CUDA

首先确定想要安装的TensorFlow版本需要哪个版本的CUDA支持。我这里安装的TF版本是1.5.0，需要CUDA 9.0的支持。

下载并安装CUDA

前往developer.nvidia.com下载CUDA 9.0。一直下一步就行。
完成之后验证一下是否已经正确安装，运行cmd：

1
nvcc  -V

如果正确安装，将会显示版本信息：

1
2
3
4
5
C:\Windows\system32>nvcc  -V
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2017 NVIDIA Corporation
Built on Fri_Sep__1_21:08:32_Central_Daylight_Time_2017
Cuda compilation tools, release 9.0, V9.0.176

安装CUDNN

同样需要确认需要哪个版本的CUDNN。我这里需要CUDNN 9.0。前往developer.nvidia.com下载指定版本，下载完成是一个压缩包，解压到C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v9.0

安装TensorFlow

1
pip install tensorflow-gpu==1.5.0

此处使用==指定版本。另外可以在pypi.python.org查询到某个包的版本信息。

测试

1
python
1
2
3
4
5
6
import tensorflow as tf
tf.__version__
sess = tf.Session()
a = tf.constant(20)
b = tf.constant(99)
print(sess.run(a + b))

运行结果：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
C:\Windows\system32>python
Python 3.5.4 (v3.5.4:3f56838, Aug  8 2017, 02:17:05) [MSC v.1900 64 bit (AMD64)] on win32
Type "help", "copyright", "credits" or "license" for more information.
>>> import tensorflow as tf
>>> tf.__version__
'1.5.0'
>>> sess = tf.Session()
2018-02-05 10:21:19.674014: I C:\tf_jenkins\workspace\rel-win\M\windows-gpu\PY\35\tensorflow\core\platform\cpu_feature_guard.cc:137] Your CPU supports instructions that this TensorFlow binary was not compiled to use: AVX AVX2
2018-02-05 10:21:20.177086: I C:\tf_jenkins\workspace\rel-win\M\windows-gpu\PY\35\tensorflow\core\common_runtime\gpu\gpu_device.cc:1105] Found device 0 with properties:
name: GeForce 940MX major: 5 minor: 0 memoryClockRate(GHz): 1.2415
pciBusID: 0000:02:00.0
totalMemory: 2.00GiB freeMemory: 1.66GiB
2018-02-05 10:21:20.177537: I C:\tf_jenkins\workspace\rel-win\M\windows-gpu\PY\35\tensorflow\core\common_runtime\gpu\gpu_device.cc:1195] Creating TensorFlow device (/device:GPU:0) -> (device: 0, name: GeForce 940MX, pci bus id: 0000:02:00.0, compute capability: 5.0)
>>> a = tf.constant(20)
>>> b = tf.constant(99)
>>> print(sess.run(a + b))
119

如何更新

1
pip install --upgrade tensorflow-gpu