谷歌JAX深度学习开发环境搭建

谷歌JAX深度学习框架要求运行在Linux系统上，我们可以在Windows11系统中，利用WSL虚拟环境搭建Ubuntu虚拟机系统，在Windows11上运行PyCharm集成开发环境，调用Ubuntu虚拟机系统上的JAX框架，对代码进行编程和调试。

（1）首先，安装Windows Subsystem for Linux（WSL）。在Windows11设置窗口上，搜索出“启用或关闭Windows功能”选项，并勾选“适用于Linux的Windows子系统”选项。重启计算机。

（2）第二，启用开发者模式。在“Windows设置”中搜索“开发人员”，打开“开发者选项”窗口，把开发人员模式打开。

 

（3）第三，升级WSL系统。

在Windows终端中以管理员身份运行如下命令：

wsl.exe --update

等待升级结束后运行如下命令：

wsl --set-default-version 2

可以通过如下命令查看WSL的版本号：

wsl --list –verbose

（4）第四，从“Microsoft Store”中安装Ubuntu。打开Microsoft Store页面，搜索Ubuntu，在搜索的结果中选择安装Ubuntu 20.04.5版本的Linux虚拟机。Microsoft Store页面打不开，自己百度一下解决。在Ubuntu 20.04.5版本的Linux虚拟机的页面上，单击“get”按钮即开始安装，安装过程很简单。

 500多M很快就装完了，设置用户名密码即可。菜单上会出来一个图标“Ubuntu 20.04.5 LTS”，点击图标即可打开Ubuntu终端。

 

（5）第五，新安装的WSL需要更新一次，打开WSL终端界面，依次输入如下操作语句：

sudo apt update

sudo apt install gcc make g++

sudo apt install build-essential

sudo apt install python3-pip

pip install --upgrade pip

（6）第六，JAX的安装和验证：

pip install -i https://pypi.tuna.tsinghua.edu.cn/simple jax==0.2.19  

pip install -i https://pypi.tuna.tsinghua.edu.cn/simple jaxlib==0.1.70

python3

这是启动WSL自带的Python命令，之后键入如下命令：

import jax.numpy as np

np.add(1.0,1.7)

如下结果说明安装无误。

 （7）第七，请读者在Windows11上直接使用PyCharm专业版，方便在编译器上直接配置WSL虚拟环境。

 

 注意：JAX相关包的安装都在Ubuntu虚拟机系统上做的，源码中缺什么包就安装什么包。

 好了。JAX框架的深度学习代码可以运行起来了。

注意：本文可以用于搭建《谷歌JAX深度学习从零开始学》那本书开发环境。