ubuntu安装JAX GPU版本

L0_L0

已于 2023-09-27 11:52:26 修改

阅读量4.3k

点赞数 3

分类专栏：随笔小记文章标签： ubuntu linux 深度学习

于 2023-09-27 00:17:09 首次发布

本文链接：https://blog.csdn.net/L0_L0/article/details/133326020

版权

随笔小记专栏收录该内容

6 篇文章

订阅专栏

ubuntu安装JAX GPU版本

首先cuda是可以向下兼容的，因此，在ubuntu22.0.4上我的cuda为12.0，cudnn为8.8；但依然可以安装jax==0.4.7 jaxlib==0.4.7+cuda11.cudnn82的版本。

0 安装前提

已经安装好cuda和对应的cudnn，以及anaconda。

从NVIDIA GPU官网中，我们可以看到当前支持的jax所对应的cuda和cudnn版本：
在这里插入图片描述

# 创建一个新的conda虚拟环境
# 在终端输入命令
conda create -n jaxEnv python=x.x

注意自己安装的python版本，创建环境指令后面的版本号要与以安装的python版本一致

# 查询创建的conda环境
conda info -e
# 激活自己刚创建的conda环境
conda activate jaxEnv

1 安装JAX GPU版本

# 可以使用上图中官方的安装命令
pip install --upgrade pip

# CUDA 12 installation
# Note: wheels only available on linux.
pip install --upgrade "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

# CUDA 11 installation
# Note: wheels only available on linux.
pip install --upgrade "jax[cuda11_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

# 或者
pip install --upgrade pip

# Installs the wheel compatible with CUDA 12 and cuDNN 8.9 or newer.
# Note: wheels only available on linux.
pip install --upgrade "jax[cuda12_local]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

# Installs the wheel compatible with CUDA 11 and cuDNN 8.6 or newer.
# Note: wheels only available on linux.
pip install --upgrade "jax[cuda11_local]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

或者可以参考google.storage，自己查询符合自己安装条件的版本
在这里插入图片描述

# 我使用的是如下命令安装成功
pip install --upgrade jax==0.4.7 jaxlib==0.4.7+cuda11.cudnn82 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

如果已经有安装的jax，切记先卸载再安装对应版本即可

# 卸载jax和jaxlib的命令
pip uninstall jax jaxlib

3 测试

测试如下文件在所创建的jaxEnv环境中是否能够得到如下类似结果

# TestJAX.py
import jax.numpy as np
from jax import random
import time

rng = random.PRNGKey(0)
x = random.uniform(rng, [5000, 5000])
st = time.time()
try:
    y = np.matmul(x, x)
except Exception:
    print("error")
print(time.time() - st)
print(y)