JAX安装

本文安装版本为:jax0.2.21、jaxlib0.1.72、tensorflow2.7
jax下载地址:https://github.com/google/jax/releases
tensorflow下载地址:https://github.com/tensorflow/tensorflow

1 安装CPU版本

pip install "jax[cpu]===0.2.21" -f https://whls.blob.core.windows.net/unstable/index.html --use-deprecated legacy-resolver

2 安装GPU版本

2.1 pip安装jax和jaxlib

(1)安装jax:pip install -v jax==0.2.21

(2)安装jaxlib:
版本对应链接:https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
下载链接:https://storage.googleapis.com/jax-releases/cuda111/jaxlib-0.1.72+cuda111-cp38-none-manylinux2010_x86_64.whl

1>在线安装:pip install jaxlib==0.1.72+cuda111 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
2>本地安装:pip install jaxlib-0.1.72+cuda111-cp38-none-manylinux2010_x86_64.whl

2.2 源码安装

(1)安装必要的python依赖:

pip install --upgrade pip -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install numpy==1.19.5 six wheel scipy==1.7.3 -i https://pypi.tuna.tsinghua.edu.cn/simple/

(2)使用修改后的tensorflow存储库从源代码构建jaxlib:

cd /public/home/xxx/data/jax/build/build.py
python build.py --enable_cuda --bazel_options=--override_repository=org_tensorflow=/public/home/xxx/data/tensorflow

(3)安装新构建的jaxlib轮子:

pip install /public/home/xxx/data/jax/build/dist/jaxlib-0.3.15-cp38-none-manylinux2014_x86_64.whl

编译成功如下图:

(4)安装jax:
在jax主目录执行:pip install -e .

你可能感兴趣的:(深度学习环境配置,tensorflow,python,深度学习)