1、直接使用pip安装即可
pip install tensorflow-gpu
2、安装 cuda
安装cuda一定要注意安装特定的版本,才能正常加载dll文件。如果不清楚需要什么版本,可以先写一个tensorflow的测试里,错误中会提示不能加载,比如:
ImportError: Could not find ‘cudart64_90.dll’
表示需要安装版本号为9的64位的cuda。
下载地址:
https://developer.nvidia.com/cuda-downloads
3、安装 cudnn
同cuda需要注意版本,可在错误中看到需要安装的版本。
下载:https://developer.nvidia.com/rdp/cudnn-download
cudnn下载完成后,把压缩包内的 bin、include、lib 复制到cuda安装目录下即可。
比如我的目录是:
C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v9.0
4、测试
import numpy import tensorflow as tf import os os.environ['CUDA_VISIBLE_DEVICES'] = "0" a = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[2, 3], name='a') b = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[3, 2], name='b') c = tf.matmul(a, b) sess = tf.Session(config=tf.ConfigProto(log_device_placement=True)) print(sess.run(c))