jax 和 jaxlib 的 cuda 版本安装 笔者花费时间才在 Ubuntu 20.04 适配上 jax 和 jaxlib 的 cuda 版本安装,以及 chex 版本。
笔者花费时间才在 Ubuntu 20.04 适配上 jax 和 jaxlib 的 cuda 版本安装,以及 chex 版本。
版本展示
本人版本展示 jax==0.4.27 ,jaxlib==0.4.27+cuda12.cudnn89,chex==0.1.86。
安装过程
cuda 以及环境变量配置过程
首先安装cuda12.4和cudnn8.9,给出链接:
cuda:https://developer.nvidia.com/cuda-toolkit-archive
cudnn:https://developer.nvidia.com/rdp/cudnn-download
安装完成cuda和cudnn后,务必将cuda写入环境变量,给出模板示例,特别是CUPTI库文件:
vim ~/.bashrc
# cuda安装位置
export CUDA_HOME=/usr/local/cuda-12.4
export PATH=$PATH:/usr/local/cuda-12.4/bin
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/cuda-12.4/lib64
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/cuda-12.4/extras/CUPTI/lib64
export LD_LIBRARY_PATH=/usr/local/cuda-12.4/targets/x86_64-linux/lib:$LD_LIBRARY_PATH
source ~/.bashrc
给出笔者配置 bashrc 图片参考:
jax 和 jaxlib 安装过程
此处首先给出官方 cuda 版本 whl 文件下载地址:
jaxlib:storage.googleapis.com/jax-releases/jax_cuda_releases.html
选择 cuda12.cudnn89 对应自己需要版本,根据自己要求选择,笔者选择为0.4.27,则执行:
pip install jax==0.4.27
# 此处jaxlib笔者为whl本地安装
pip install jaxlib-0.4.27+cuda12.cudnn89-cp39-cp39-manylinux2014_x86_64.whl
切记笔者 chex==0.1.86,chex 对 jax 也有要求,容易遇到 bug ,全是泪与恨。
验证是否 GPU 可用
在终端内验证:
>>> import jax
>>> print(jax.devices())
[cuda(id=0)]
正确输出 cuda id 即为成功。
参考(Reference)
jax 和 jaxlib 的 cuda 版本安装 笔者花费时间才在 Ubuntu 20.04 适配上 jax 和 jaxlib 的 cuda 版本安装,以及 chex 版本。jax 和 jaxlib 的 cuda 版本安装 笔者花费时间才在 Ubuntu 20.04 适配上 jax 和 jaxlib 的 cuda 版本安装,以及 chex 版本。jax 和 jaxlib 的 cuda 版本安装 笔者花费时间才在 Ubuntu 20.04 适配上 jax 和 jaxlib 的 cuda 版本安装,以及 chex 版本。jax 和 jaxlib 的 cuda 版本安装 笔者花费时间才在 Ubuntu 20.04 适配上 jax 和 jaxlib 的 cuda 版本安装,以及 chex 版本。jax 和 jaxlib 的 cuda 版本安装 笔者花费时间才在 Ubuntu 20.04 适配上 jax 和 jaxlib 的 cuda 版本安装,以及 chex 版本。jax 和 jaxlib 的 cuda 版本安装 笔者花费时间才在 Ubuntu 20.04 适配上 jax 和 jaxlib 的 cuda 版本安装,以及 chex 版本。jax 和 jaxlib 的 cuda 版本安装 笔者花费时间才在 Ubuntu 20.04 适配上 jax 和 jaxlib 的 cuda 版本安装,以及 chex 版本。jax 和 jaxlib 的 cuda 版本安装 笔者花费时间才在 Ubuntu 20.04 适配上 jax 和 jaxlib 的 cuda 版本安装,以及 chex 版本。jax 和 jaxlib 的 cuda 版本安装 笔者花费时间才在 Ubuntu 20.04 适配上 jax 和 jaxlib 的 cuda 版本安装,以及 chex 版本。