Tensorflow使用Gpu训练模型,并不是简单的使用pip install tensorflow-gpu就安装完成了。还需要在电脑上安装cuda和cudnn。本教程介绍如何在window10系统上安装tensorflow-gpu。
1.查看tensorflow对应的cuda版本
https://tensorflow.google.cn/install/source_windows
我安装的是tensorflow_gpu-1.4.0,所以下载cuda版本为v8,cudnn版本为v6
tensorflow_gpu-1.4.0 | 3.5-3.6 | MSVC 2015 update 3 | Cmake v3.6.3 | 6 | 8 |
2.下载并安装cuda
https://developer.nvidia.cn/cuda-toolkit-archive
下载后双击安装,查看是否安装成功:
控制台输入:nvcc -V
3.下载并安装cudnn
首先访问:https://developer.nvidia.cn/cudnn,进入如下界面,点击下载需要注册账号,并做完一个问卷调查,重新登陆后可以下载。
下载完成后,解压如下:
然后将每个文件夹中的文件复制到cuda对应的文件夹中去:
4.安装tensorflow-gpu
pip install tensorflow-gpu==1.4
5.重启电脑
查看是否安装成功,
print(tf.test.is_gpu_available())