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
文件,根据提示安装即可。