移动端 3060 6G 配置 TensorFlow 和 Torch

前置条件

硬件条件

  R7-6800H + RTX3060（移动端）

检查方法

  检测 Tensorflow 是否检测得到 GPU，运行以下 Python 代码

import tensorflow as tf
gpus = tf.config.experimental.list_physical_devices(device_type='GPU')
cpus = tf.config.experimental.list_physical_devices(device_type='CPU')
print(gpus, cpus)

  如果运行结果为

[] [PhysicalDevice(name='/physical_device:CPU:0', device_type='CPU')]

则说明 TF 没有检测到 GPU。

  检测 Pytorch 是否检测得到 GPU，运行以下 Python 代码

import torch
print(torch.cuda.is_available())  # 判断是否可以使用gpu计算
print(torch.cuda.device_count())  # 显示gpu数量

  如果运行结果为

False
0

则说明 Pytorch 没有检测到 GPU。

解决方法

  我是看了别的大佬的文章整理的^{1 2}，作为笔记总结供自己日后参考，所以以下我只会直接给出对应的版本和安装指令。

  所有的包如果之前有下载过不同版本的，最好是先卸载再安装。

GPU 驱动

CUDA==10.1

  用 conda 下载安装，pip 不行，具体原因未知

conda install cudatoolkit=10.1

cudnn==7.6.5

  同样只能用 conda 下载安装

conda install cudnn=7.6

Tensorflow

tensorflow==2.3.0

  这里可以稍微注意一下，用 pip 安装 tensorflow 之前，最好确保 Python 版本在 3.5~3.8 之间³，否则可能会出现报错。

  这里就只能用 pip 下载安装了

pip install tensorflow-gpu==2.3.0

Pytorch

torch==1.8.1+cu101

  Pytorch 官网其实有提供官方的版本对齐文档⁴，由于上面我们的 CUDA 版本为 10.1，所以我直接选择了最新的对应 CUDA 10.1 的 torch 版本。

  这里可以在网页上 Ctrl+F 调出搜索框，然后直接搜索你要的版本。

# CUDA 10.1
pip install torch==1.8.1+cu101 torchvision==0.9.1+cu101 torchaudio==0.8.1 -f https://download.pytorch.org/whl/torch_stable.html

结果

  再次检查，Tensorflow 对应的输出变为

[PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU')] [PhysicalDevice(name='/physical_device:CPU:0', device_type='CPU')]

Pytorch 对应的输出变为

True
1

则代表都检测成功了。