如何在Google Colab中将CUDA升级至11.2、CUDNN升级至8.1

原有脚本失效核心原因 #

• 安装顺序错误：你在添加NVIDIA软件源之前就执行了CUDA toolkit安装，此时默认系统源的CUDA版本不符合要求，相当于提前安装了错误版本
• 系统版本不匹配：你使用的是Ubuntu16.04的NVIDIA源，当前Google Colab默认运行环境为Ubuntu20.04，源不匹配会导致安装包依赖报错、版本安装失败
• 软链接未更新：安装完CUDA 11.2后没有修改/usr/local/cuda默认软链接指向，系统会优先调用原有旧版本CUDA
• 环境变量生效范围有限：你只修改了当前Python进程的环境变量，shell命令的全局环境变量没有同步，所以执行nvcc --version等命令时还是返回旧版本信息

修改后的可用安装脚本 #

# 先清理原有CUDA残留包和文件
!sudo apt-get --purge remove "*cublas*" "*cufft*" "*curand*" "*cusolver*" "*cusparse*" "*npp*" "*nvjpeg*" "cuda*" "nsight*" -y
!sudo apt-get autoremove -y
!sudo apt-get autoclean
!sudo rm -rf /usr/local/cuda*

# 添加适配Ubuntu20.04的NVIDIA官方CUDA源
!wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/cuda-keyring_1.0-1_all.deb
!sudo dpkg -i cuda-keyring_1.0-1_all.deb
!sudo apt-get update

# 安装CUDA 11.2全量组件
!sudo apt-get install -y cuda-11-2

# 更新默认CUDA软链接指向11.2版本
!sudo ln -sf /usr/local/cuda-11.2 /usr/local/cuda

# 同步配置全局环境变量，同时对shell和Python进程生效
!echo 'export PATH=/usr/local/cuda/bin:$PATH' >> ~/.bashrc
!echo 'export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH' >> ~/.bashrc
%env PATH=/usr/local/cuda/bin:/usr/local/nvidia/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:/tools/node/bin:/tools/google-cloud-sdk/bin:/opt/bin
%env LD_LIBRARY_PATH=/usr/local/cuda/lib64:/usr/lib64-nvidia

# 安装CUDNN 8.1，需提前把cudnn安装包上传到对应Google Drive路径
!cp "/content/drive/My Drive/cudnn-11.2-linux-x64-v8.1.0.0.tgz" "/content/cudnn-11.2.tgz"
!tar -zxvf cudnn-11.2.tgz
!sudo cp cuda/include/cudnn*.h /usr/local/cuda/include
!sudo cp cuda/lib64/libcudnn* /usr/local/cuda/lib64
!sudo chmod a+r /usr/local/cuda/include/cudnn*.h /usr/local/cuda/lib64/libcudnn*
!sudo ldconfig

版本验证注意事项 #

脚本执行完成后，必须先重启Colab运行时（点击菜单栏「代码执行程序」→「重启运行时」），再执行版本校验代码，否则环境变量不会完全生效。
推荐用以下代码验证：

# 校验CUDA版本
!nvcc --version
# 校验CUDNN版本
!cat /usr/local/cuda/include/cudnn_version.h | grep CUDNN_MAJOR -A 2
# 校验TensorFlow是否能正常识别GPU
import tensorflow as tf
print(tf.config.list_physical_devices('GPU'))

内容的提问来源于stack exchange，提问作者deepesh mhatre