Tensorflow-gpu2

发表于 分类于 阅读次数： Waline： 本文字数： 664 阅读时长 ≈ 1 分钟
浅浅做一个踩坑记录

前言

装好 Tensorflow 2.3.0 之后，训练一个 MNIST 集花了7分钟左右，于是想安装 GPU 版本的. 于是就开始踩坑了

选定版本

安装 Tensorflow-gpu 需要安装相应的 CUDA 和 cuDNN 版本。这里应该先看自己的显卡所支持的 CUDA 版本，再确定 Tensorflow 版本。通过 NIVIDIA-系统信息-组件 查看

驱动支持的CUDA版本

我的显卡是 RTX 3060 版本比较高，之前安装的 2.3.0 版本不知道为什么就在编译模型，和构建模型的时候卡死。所以选择 Tensorflow-gpu=2.4.0 。

安装

打开cmd, 激活虚拟环境

1
conda active ts

选择阿里云镜像源，默认源无该版本且速度慢

1
pip install Tensorflow-gpu=2.4.0 -i https://mirrors.aliyun.com/pypi/simple/

听说这里应该会自动安装对应的 CUDA 和 cuDNN，然而我并没用，手动安装

1
pip install cudnn==8.0 -i https://mirrors.aliyun.com/pypi/simple/

这里会自动安装 cudatookit ，接着安装 keras

1
pip install keras==2.4.3

是否检测到GPU

1
2
import tensorflow as tf
tf.config.list_physical_devices('GPU')

输出 [PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU')] 则成功。