深度学习 SOP: 教你在Conda中搭建任意版本TensorFlow GPU环境

4 阅读4分钟

背景

在深度学习环境配置中,尤其是TensorFlow的GPU环境,对各个依赖的版本非常苛刻,本文旨在探索一种持续可用的SOP,方便配置任何版本的TensorFlowGPU环境。

在这个SOP中不会指定任何一个版本,而是指导你如何 根据你的需要 的版本取去查询对应的依赖版本!这样能让这个SOP持续可用!

准备工作

在开始安装之前,我们需要了解三个关键版本,这决定了环境是否能兼容运行。

  • Tensorflow版本
  • CUDA版本
  • cuDNN版本

不用担心复杂,官方已经整理好了依赖关系,你需要做的就只有确保你的电脑已经正确安装了GPU驱动,这可以通过命令nvidia-smi来进行查看,在命令结果的右上角,你可以看到你驱动所支持的最大CUDA版本

*最小版本其实也是有限制的,这里我们先不关注,文章最后有说明。


具体步骤

1. 前往TensorFlow官网查看目标版本所需依赖

经过测试的构建配置(注意一定要语言一定选英文,中文版的很久没更新了!!)上方这个链接里TensorFlow官方给出了不同TensorFlow版本所兼容的CUDA和cuDNN版本。

假如现在团队要求安装Tensorflow 2.3.0版本。我们就找到这一行:

版本Python 版本编译器构建工具cuDNNCUDA
tensorflow-2.3.03.5-3.8GCC 7.3.1Bazel 3.1.07.610.1

取出我们所需的三个关键版本

  • Python版本:3.8
  • CUDA版本:10.1
  • cuDNN:7.6

当然也可以根据你的需要选择版本,主要就是取到这三个关键版本,Python版本一般选支持范围最新的就可以。


2. 创建conda环境并激活

替换tf23为你想创建环境名称、替换3.8为Python版本。

conda create -n tf23 python=3.8 -y
conda activate tf23

3. 安装环境独立的CUDA/cuDNN

替换10.17.6为你需要的CUDA/cuDNN版本。

conda install -y cudatoolkit=10.1 cudnn=7.6

有些版本提示找不到,可以选用社区源

conda install -y cudatoolkit=10.1 cudnn=7.6 -c conda-forge

*如果慢且长时间安装不成功,请参考附录的解决方案。


4. 安装 TensorFlow

替换2.3.0为你需要的TensorFlow版本

pip install "tensorflow==2.3.0"

这里建议直接装 tensorflow==2.3.0。原因是 从 TF 2.1 起,tensorflow 与 tensorflow-gpu 本质上是同一个包(后来官方也明确让大家装 tensorflow)。

*如果感觉慢,或者报错433。可以尝试临时使用国内源(这里是清华源)。

pip install "tensorflow==2.3.0" -i https://pypi.tuna.tsinghua.edu.cn/simple

附录:

如果在第三步中感觉慢,可以尝试临时使用国内源(这里是清华源)。

conda install -y \
  cudatoolkit=10.1 \
  cudnn=7.6 \
  -c https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge \
  --override-channels

或者长期更换conda源。清华源Conda换源教程 首先查看现有源。

conda config --show-sources

一般是只有一个default,通过这个命令添加自定义源(这里是清华源)。

conda config --set custom_channels.auto https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/

再次查看现有源,可以看到已经添加了清华源,此时可以再次尝试安装。

conda config --show-sources

如果出现错误,可以用这个命令移除。

conda config --remove channels 源名称

总结

本文的主要场景是针对团队合作、项目复现、导师要求等情况下需要安装【指定版本】的SOP。当然安装最新版本也可以按照这个流程自行调整。

番外——检查当前GPU所支持TensorFlow版本范围

给定一台 GPU,判断:

  • 是否支持某个 TensorFlow 版本
  • 支持 GPU 的 最老/最新 TensorFlow 版本

1. 确定 GPU 的 Compute Capability(CC)

NVIDIA 官方表 developer.nvidia.com/cuda-gpus

  1. 命令行运行nvidia-smi
  2. 拿到 GPU 型号(如 RTX 3060)
  3. 在上面的 NVIDIA 页面查到 Compute Capability(如 8.6)

2. 由 Compute Capability 推导最低 CUDA 版本

NVIDIA CUDA Release Notes docs.nvidia.com/cuda/ CUDA release notes 中会标注 “Added support for compute capability X.Y 这里也有一个常用的

# Compute Capability → 最低 CUDA(参考)
CC 5.x (Maxwell)     → CUDA ≥ 7.0
CC 6.x (Pascal)      → CUDA ≥ 8.0
CC 7.0 (Volta)       → CUDA ≥ 9.0
CC 7.5 (Turing)      → CUDA ≥ 10.0
CC 8.x (Ampere)      → CUDA ≥ 11.1
CC 9.x (Hopper)      → CUDA ≥ 12.0

3. 查询GPU驱动支持最大CUDA版本

  1. 命令行运行nvidia-smi
  2. 右上角CUDA Version即驱动支持最大版本

4. 查 TensorFlow 版本使用的 CUDA 版本

TensorFlow 官方安装文档(历史版本)www.tensorflow.org/install/sou… 查每个 TF 版本对应的:

  • CUDA version
  • cuDNN version

5. 兼容性判定公式

GPU 要求的最低 CUDA 版本TF 使用的 CUDA版本GPU 驱动支持的最大 CUDA 版本

5. 得到TensorFlow可使用范围

在表格中找到 CUDA 版本对应的 TensorFlow版本


示例(RTX 3060)

  • GPU:RTX 3060
  • CC:8.6(NVIDIA CUDA GPU 表)
  • 最低 CUDA:11.1(CUDA 11.1 release notes)
  • TF 2.3 → CUDA 10.1 ❌
  • TF 2.4 → CUDA 11.0 ❌
  • TF 2.5 → CUDA 11.2 ✅

结论:RTX 3060 最老支持 TF ≈ 2.5+