2019 年写了一篇《一条命令构建 Tensorflow 开发环境》,使用 conda install tensorflow-gpu
可以自动安装 cudatoolkit
、cudnn
、tensorflow-gpu
,非常方便入门同学搭建 Tensorflow GPU 开发环境,不折腾。
2023 年,tensorflow-gpu
已经合并到 tensorflow
,Apple 的 M1 / M2 系列芯片也支持 Tensorflow / Pytorch 加速了,重新写一篇记录一下如何快速搭建 Tensorflow GPU 开发环境,方便刚入门的同学。
以下的指南适用于:tensorflow 2.12
、python 3.10
Miniconda 是 Anaconda 的精简版本,只保留了 Conda 和 Python,比 Anaconda 简洁。安装 Miniconda 可以使用“清华镜像“。
Ubuntu 上最简单的 Nivida GPU 驱动安装方法是通过 apt 的第三方源 PPA。
# 添加 ppa 显卡驱动源
sudo add-apt-repository ppa:graphics-drivers/ppa
sudo apt update
# 查看可用的驱动版本
ubuntu-drivers devices
# 安装指定版本的驱动
sudo apt install nvidia-driver-530
# 重启后,检查是否安装成功
nvidia-smi
Tips:
有的教程指出:使用sudo ubuntu-drivers autoinstall
可以自动选择合适的版本进行安装。这种方式在 2023 年已经被弃用了。执行ubuntu-drivers --help
将看到以下提示:Commands: autoinstall Deprecated, please use "install" instead
按照官方的指南安装 cudatoolkit 和 cudnn 很费劲,使用 conda 来安装就轻松多了。
安装时,需注意预编译的 tensorflow 和 cudnn、cudatoolkit 之间的版本关系。 对 tensorflow 2.12.0
,官方预编译版本要求:cudnn 8.6+
,cudatoolkit 11.8
。conda 上 cudatoolkit 和 cudnn 的版本并不完整,经测试可以匹配的版本是:cudnn 8.8
和 cudatoolkit 11.8
。
使用 conda 安装 cudatoolkit 时不会安装 nvcc 和 cupti,会导致 XLA 加速库调用失败,显示错误信息:InternalError: libdevice not found at ./libdevice.10.bc
。可以通过安装 cuda-nvcc
及 cuda-cupti
并设置 XLA_FLAGS
来解决这个问题。
使用以下命令安装以上各个包:
conda install -c conda-forge cudatoolkit=11.8 cudnn=8.8
conda install -c nvidia cuda-nvcc=11.8 cuda-cupti=11.8
Tips: Tensorflow与Python、CUDA、cuDNN的版本对应表
必须正确设置 LD_LIBRARY_PATH
指向 cudnn library 的目录;并设置 XLA_FLAGS
export LD_LIBRARY_PATH=$CONDA_PREFIX/lib
export XLA_FLAGS=--xla_gpu_cuda_data_dir=$CONDA_PREFIX
使用以下命令可以让 conda 在激活环境时自动设置该环境变量。
# 该命令对当前环境有效,否则需使用 -n 参数指定环境名
conda env config vars set LD_LIBRARY_PATH=$CONDA_PREFIX/lib
conda env config vars set XLA_FLAGS=--xla_gpu_cuda_data_dir=$CONDA_PREFIX
pip install tensorflow
import tensorflow as tf
tf.test.is_gpu_available()
tf.config.list_physical_devices('GPU')
安装下面两个包就可以使用 Apple M1 / M2 上的 GPU 进行加速了:
pip install tensorflow-macos
pip install tensorflow-metal
Tips: Apple 官方指南 一直没有更新,已经过时了!2023 年,对于
tensorflow-macos 2.12.0
来说,无需安装conda install -c apple tensorflow-deps
,安装了反而会出错!只有tensorflow-macos <= 2.10
的版本才需要安装。
pytorch 就简单多了,无论是什么平台,一条指令就自动安装所有的包,无需单独安装 GPU 的支持包:
pip install torch