Ubuntu + RTX 30 series 安装 TensorFlow-GPU

原创

已于 2022-05-18 21:51:40 修改 · 1k 阅读

4 ·

CC 4.0 BY-SA版权

文章标签：

#tensorflow #ubuntu

于 2022-05-18 15:02:02 首次发布

本文介绍了在Ubuntu系统上，针对RTX 30系列显卡安装TensorFlow-GPU的过程中遇到的问题及解决方案。主要内容包括选择适合RTX 30系列的TensorFlow版本，探讨了Docker和Anaconda两种环境隔离方法，以及详细步骤如安装NVIDIA驱动、Docker、下载TensorFlow镜像和创建新环境。

Ubuntu + RTX 30 series 安装 TensorFlow-GPU

在实验室服务器上安装TensorFlow遇到挺多坑，包括30系显卡适用的TensorFlow版本，以及如何在公用环境中配置自己的环境时不影响到别人，所以记录一下这些问题和解决方案。

GeForce RTX 30 series 显卡匹配的 TensorFlow 版本

RTX 30系显卡采用了新的Ampere架构，支持CUDA 11.1及以上版本。根据TensorFlow官网提供的对应关系，官方支持的版本中需要tensorflow-2.5.0及以上才适配到CUDA 11.2，也就是官方支持的版本完全不适配tensorflow-1.x。

但是为了加速创新，使研发人员能快速构建环境，NVIDIA官方提供了一个镜像仓库NGC，提供多种深度学习框架的docker镜像。其中，Google Brain Team 维护了TensorFlow的镜像更新，支持最新版本的CUDA与TensorFlow适配。不同镜像版本支持的操作系统、CUDA版本、TensorFLow版本等在手册中进行了说明，每个Release版本包含一个最新的 TensorFlow 1.x 版本和一个 2.x 版本。

环境的独立性

每个 TensorFlow 版本都需要不同的 CUDA 版本，因此 CUDA 不应该安装在基础环境中，虽然我在上篇博客里写了主机环境中 CUDA 版本的切换，但这种方式并不利于维护，对以后的使用者会造成很多麻烦。因此，推荐以下两种配置方式。

第一种是使用 Docker 安装上述 NGC 容器环境，配置和使用简单，主机只需要具备显卡驱动即可。缺点是镜像比较大，如果需要多个版本的切换，比较冗余。

第二种是在 Anaconda 里安装 CUDA ，在虚拟环境中使用conda install cudatoolkit=11.2.0命令安装，不改变基础环境。但这种方式可能需要自己研究一下版本的匹配关系。

准备工作

安装NVIDIA显卡驱动：

选择显卡型号和操作系统（以RTX 3090/Ubuntu 20.04为例），点击Search获取显卡驱动下载链接，然后运行.run文件，根据提示安装即可。

wget https://us.download.nvidia.com/XFree86/Linux-x86_64/510.73.05

最低0.47元/天解锁文章