Ubuntu系统 TensorFlow使用GPU训练

前言

本教程的前提是在ubuntu系统内，已成功安装NVIDIA显卡驱动，若没有安装NVIDIA显卡驱动点这里

本人亲测可用配置：

ubuntu 20.04 + GTX 1080Ti + TensorFlow-gpu 2.3.1 + CUDA 10.1 + CUdnn 7.6.0 + python 3.8.10

一、安装CUDA

1.1 确定CUDA版本

首先，不同TensorFlow版本对应的CUDA版本也是不同的。如果你安装最新或最高支持的版本，但tensorflow不支持，也没啥用。TensorFlow官网给出的测试配置在这

1.2 下载CUDA文件

下载地址

选择适合自己的CUDA版本，runfile文件下载，并按照安装命令安装

1.3 安装CUDA（可实现多版本cuda切换）

CUDA 10.0安装方法教程参考

CUDA 10.1及以上安装的方法教程参考

二、安装CUDNN

CUDNN的安装要根据所装的CUDA版本进行对应，下载前需要注册NVIDIA账号，CUDNN各版本下载官网

CUDNN8.0以下版本的安装方法，参考

CUDNN8.0及以上版本的安装方法，参考

三、环境配置

打开终端，输入

sudo gedit ~/.bashrc
# 添加以下内容
# 支持多个CUDA版本切换,要改版本就修改/usr/local/cuda的软连接就行了
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/cuda/lib64
export PATH=$PATH:/usr/local/cuda/bin
export CUDA_HOME=$CUDA_HOME:/usr/local/cuda

如果使用的是conda虚拟环境，一定要在环境内安装对应cudatoolkit版本，才能正常使用

命令如下：

# 激活自己的环境
conda activate tf2rl
# cuda安装哪个版本，就安装对应cudatoolkit版本
# 本例使用cuda 10.1
conda install cudatoolkit==10.1

到此，环境配置完成后，必须重启电脑，不然检测还是False！

四、环境测试

检测CUDA和CUDNN是否安装完成

检测cuda版本 ： nvcc -V
查看位置 ： which nvcc
查看NVIDIA动态使用情况： watch -n 1 nvidia-smi
cuda 版本 ： cat /usr/local/cuda/version.txt
检测cudnn 版本 ： cat /usr/local/cuda/include/cudnn.h | grep CUDNN_MAJOR -A 2
NVIDIA 驱动版本 ： cat /proc/driver/nvidia/version

查看TensorFlow gpu是否可用

# 在自己环境内的python
import tensorflow as tf
tf.test.is_gpu_available()

当显示True时，表示TensorFlow和GPU连接成功，大功告成！！