IT评测·应用市场-qidao123.com

标题: jax 和 jaxlib 的 cuda 版本安装 [打印本页]

作者: 石小疯 时间: 2024-10-20 09:23
标题: jax 和 jaxlib 的 cuda 版本安装
笔者花费时间才在 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: 库安装和GPU使用，解决不能识别gpu问题_jax找不到gpu-CSDN博客
JAX安装过程中碰到的各种坑-CSDN博客

免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！更多信息从访问主页：qidao123.com:ToB企服之家，中国第一个企服评测及商务社交产业平台。

欢迎光临 IT评测·应用市场-qidao123.com (https://dis.qidao123.com/)