flax与jax、jaxlib的安装顺序和版本匹配


前言

在复现代码时需要用到这三个库,在安装的时候遇到了一些问题,现将解决措施记录如下。


一、安装顺序及方法

顺序是:jaxlib→jax→flax,安装之后记得检查能否使用GPU
jaxlib安装方法:
(1)Windows:https://whls.blob.core.windows.net/unstable/index.html
(2)Linux或Mac:直接在https://pypi.org/上下载安装包并安装
jax安装:https://blog.csdn.net/Papageno_Xue/article/details/125754893
flax也是直接在pypi上下载安装包并安装。

二、版本匹配

1.基础环境

这三个库是在TensorFlow环境当中安装的,因此版本首先需要与tf版本匹配。基础环境的版本如下:
ml-collections 0.1.1
tensorflow-gan 2.1.0
tensorflow-gpu 2.9.0
tensorflow-hub 0.12.0
tensorflow-io 0.28.0
tensorflow-addons 0.17.1
tensorflow-datasets 4.3.0
tensorflow-probability 0.16.0
tensorboard 2.9.1
cuda 11.2,cudnn 8.2.0


2.三个库的版本匹配

jaxlib 0.3.14+cuda11.cudnn82;jax 0.3.14;flax 0.5.3。
这里介绍一个很有用的工具pipdeptree,如果需要查看某个库的依赖库有哪些,包括对应的版本,可以用 pipdeptree -p 库名 来查看。

总结

本文介绍了flax与jax、jaxlib的安装顺序和版本匹配,希望对大家有帮助。

猜你喜欢

转载自blog.csdn.net/syfei0518/article/details/131253622