通过 requirements.txt 使用 CUDA 安装 JAX

Question

我使用官方 JAX Github的安装说明安装了 JAX[cuda]： pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html效果很好。 现在我想将其包装到 requirements.txt 中以保存我的环境。

我已经在我的 requirements.txt 中尝试--pre --extra-index-url https://storage.googleapis.com/jax-releases/jax_cuda_releases.html "jax[cuda]"但是pip install -r requirements.txt总是安装 CPU 版本，而不是 GPU 版本。 是否有任何技巧可以强制要求文件安装 JAX 的 GPU / cuda 版本？

Answer 1

在pip requirements.txt文件中，每个全局选项都应该写在自己的单独行中：

--pre

--find-links https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

jax[cuda]

参考

• https://pip.pypa.io/en/stable/reference/requirements-file-format/#supported-options