Ubuntu + RTX 30 series 安装 TensorFlow-GPU
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/NVIDIA-Linux-x86_64-510.73.05.run
./NVIDIA-Linux-x86_64-510.73.05.run
通过 Docker 安装 TensorFlow
- 安装 Docker
sudo apt-get update
sudo apt-get install ca-certificates curl gnupg lsb-release
# 添加 Docker 的官方 GPG 密钥
curl -fsSL https://download.docker.com/linux/ubuntu/gpg | sudo gpg --dearmor -o /usr/share/keyrings/docker-archive-keyring.gpg
echo "deb [arch=$(dpkg --print-architecture) signed-by=/usr/share/keyrings/docker-archive-keyring.gpg] https://download.docker.com/linux/ubuntu $(lsb_release -cs) stable" | sudo tee /etc/apt/sources.list.d/docker.list > /dev/null
sudo apt-get update
sudo apt-get install docker-ce docker-ce-cli containerd.io docker-compose-plugin
# 验证是否已经正确安装
sudo docker run hello-world
- 安装NVIDIA Container Toolkit
distribution=$(. /etc/os-release;echo $ID$VERSION_ID) \
&& curl -s -L https://nvidia.github.io/libnvidia-container/gpgkey | sudo apt-key add - \
&& curl -s -L https://nvidia.github.io/libnvidia-container/$distribution/libnvidia-container.list | sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list
sudo apt-get update
sudo apt-get install -y nvidia-container-toolkit
测试安装
sudo docker pull docker.io/nvidia/cuda:11.0.3-base-ubuntu20.04
sudo docker run --rm --gpus 0 -t docker.io/nvidia/cuda:11.0.3-base-ubuntu20.04 nvidia-smi
- 下载镜像
根据自己的需求选择Release 版本，然后在对应的版本后点击Pull Tag，其中 tf1 对应 1.x 版本， tf2 对应 2.x 版本。例如选择 22.04-tf1-py3，并点击按钮后，pull 命令被写入剪贴板如下：
docker pull nvcr.io/nvidia/tensorflow:22.04-tf2-py3
拉取镜像即可。
- 运行容器
# -d 后台运行
# -it 交互式
# -v dir1:dir2 目录挂载，dir1为主机目录，dir2为容器内部目录
# -p port1:port2 端口映射，port1为主机端口，port2为容器端口，8888为默认为jupyter配置的端口
# --name 容器名称
sudo docker run --gpus all -d -it -v /home/zhy/LSTM:/workspace/LSTM -p 8888:8888 --name=tf1 nvcr.io/nvidia/tensorflow:22.04-tf1-py3
# 进入容器内部工作区
sudo docker exec -it tf1 /bin/bash
通过 Anaconda 安装 TensorFlow 2.x
直接通过conda install tensorflow-gpu=2.6可以安装，但是依赖的版本可能存在问题。
- 确定安装版本
conda config --add channels conda-forge
conda search tensorflow-gpu
比如搜索结果中如下：
tensorflow-gpu 2.7.0 cuda112py38h0bbbad9_0 conda-forge
这个 build 版本中对应的 CUDA 版本为11.2， python 版本为3.8， 安装时除了要限定版本号，还要指定build编号。
根据cudnn下载页面，可以选择 8.1-8.2的 cudnn。
- 安装新环境
conda create -n tf2-gpu python=3.8
conda activate tf2-gpu
conda install cudatoolkit=11.2.0
conda install cudnn=8.2.1
conda install tensorflow-gpu=2.7.0=cuda112py38h0bbbad9_0