确认显卡驱动正确安装:

(notebook) [wuhf@aps ~]$ nvidia-smi
Thu Aug 20 18:07:33 2020
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 430.50       Driver Version: 430.50       CUDA Version: 10.1     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|===============================+======================+======================|
|   0  GeForce GTX 1080    Off  | 00000000:82:00.0 Off |                  N/A |
| 27%   29C    P8     6W / 180W |    113MiB /  8119MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+
| Processes:                                                       GPU Memory |
|  GPU       PID   Type   Process name                             Usage      |
|=============================================================================|
|    0     37575      C   ...uhf/miniconda3/envs/notebook/bin/python   103MiB |
+-----------------------------------------------------------------------------+

CUDA版本和Tensorflow版本有对应关系,TF2.0可以使用CUDA 10.1,安装TF2.0版本,查看conda 源中的TF :

(notebook) [wuhf@aps ~]$ conda search tensorflow |grep 2.0.0
tensorflow                     2.0.0 eigen_py27hec4e49e_0  pkgs/main
tensorflow                     2.0.0 eigen_py36ha83d16c_0  pkgs/main
tensorflow                     2.0.0 eigen_py37hce6be7f_0  pkgs/main
tensorflow                     2.0.0 gpu_py27hb041a2f_0  pkgs/main
tensorflow                     2.0.0 gpu_py36h6b29c10_0  pkgs/main
tensorflow                     2.0.0 gpu_py37h768510d_0  pkgs/main
tensorflow                     2.0.0 mkl_py27h68eb67f_0  pkgs/main
tensorflow                     2.0.0 mkl_py36hef7ec59_0  pkgs/main
tensorflow                     2.0.0 mkl_py37h66b46cc_0  pkgs/mai

一定要安装 gpu的build,指定build安装方法:

conda install {project}={version}={build}

执行命令:

conda install tensorflow=2.0.0=gpu_py36h6b29c10_0

然后来执行python代码测试TF是否正常:

import tensorflow as tf
tf.test.is_gpu_available()
tf.test.gpu_device_name()

输出:

Type "help", "copyright", "credits" or "license" for more information.
>>> import tensorflow as tf
>>> tf.test.is_gpu_available()
True
>>> tf.test.gpu_device_name()
2020-08-20 18:16:13.857330: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1618] Found device 0 with properties:
'/device:GPU:0'
>>>

如果安装很慢,可使用国内源,在用户目录下新建.condarc文件,内容如下:

channels:
  - defaults
show_channel_urls: true
channel_alias: https://mirrors.tuna.tsinghua.edu.cn/anaconda
default_channels:
  - https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main
  - https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/free
  - https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/r
  - https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/pro
  - https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/msys2
custom_channels:
  conda-forge: https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud
  msys2: https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud
  bioconda: https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud
  menpo: https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud
  pytorch: https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud
  simpleitk: https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud

使用pip安装

不管是gpu还是cpu版本的tf,安装文件都会复制到{python}/lib/python3.7/site-packages/tensorflow 这样后安装就会覆盖新安装的。

关键点总结:

  1. tensorflow一定要安装带gpu的build
  2. tensorflow-gpu这个不用安装
  3. tensorflow的版本要和cuda版本匹配,可以去官网查看对应关系

资料:
Tensorflow与CUDA之间的关系: https://tensorflow.google.cn/install/gpu