开始使用 JAX！ :computer: :zap:

本仓库的目标是让初学者更容易上手 JAX、Flax 和 Haiku！

JAX 生态系统正逐渐成为 PyTorch 和 TensorFlow 的一个越来越受欢迎的替代方案。 :sunglasses:

注：我在这里只会推荐那些我个人分析过并认为有用的内容。 如果你想要一份更全面的列表，可以查看 awesome-jax 仓库。

目录

• 使用 JAX 进行机器学习
  • 教程 #1：从零到英雄
  • 教程 #2：从英雄到超级英雄+
  • 教程 #3：用纯 JAX 从头开始编写神经网络
  • 教程 #4：Flax 从零到英雄
  • 教程 #5：Haiku 从零到英雄（即将推出）
• 其他有用的 JAX 资源

我的 JAX 机器学习教程

使用笔记本的提示：直接在 Google Colab 中打开笔记本即可 （你会在 Jupyter 文件顶部看到一个按钮，它会将你引导至 Colab）。 这样你就可以避免设置 Python 环境了！（这对我尤其方便，因为我使用的是目前仍不被支持的 Windows 系统）

教程 #1：从零到英雄

在这个视频中，我们从基础知识开始，然后逐步深入探讨 jit、grad、vmap 以及 JAX 的各种独特之处。

YouTube 视频（教程 #1）
配套 Jupyter 笔记本

教程 #2：从英雄到超级英雄+

在这个视频中，我们将学习训练机器学习模型（如神经网络）所需的全部附加组件，以便在多台机器上进行分布式训练！ 我们将训练一个简单的 MLP 模型，甚至会在 8 个 TPU 核心上训练一个机器学习模型！

YouTube 视频（教程 #2）
配套 Jupyter 笔记本

教程 #3：从零开始构建神经网络

观看我如何从头开始编写一个神经网络！ :partying_face: 这是 JAX 教程系列的第三集。

在这个视频中，我构建了一个 MLP，并使用 PyTorch 的数据加载器将其作为分类器在 MNIST 数据集上进行训练（当然也可以轻松使用更复杂的数据集）——所有这些都完全基于“纯”JAX 实现，未使用 Flax、Haiku 或 Optax。

随后我还进行了额外的分析：

• 可视化 MLP 学习到的权重
• 使用 t-SNE 可视化一批图像的嵌入
• 最后，我分析了我们的网络中是否有过多的 ReLU 神经元处于“死亡”状态。

YouTube 视频（教程 #3）
配套 Jupyter 笔记本（注：我很快会对其进行重构，但这里先提供原始版本）

---

教程 #4：使用 Flax 进行机器学习——从零到英雄

在这个视频中，我将介绍你开始使用 Flax 所需的一切知识！

我们将涵盖 init、apply、TrainState 等内容，以及其他一些独特的概念，比如 mutable 和 rngs 关键字的使用。

YouTube 视频（教程 #4）
配套 Jupyter 笔记本

---

教程 #5（即将推出）：使用 Haiku 进行机器学习——从零到英雄

待办事项

其他有用的内容

除了 官方文档 之外，以下资源也对我有所帮助。

视频

• JAX 简介（提供非常高层次的概述）
• JAX：加速机器学习研究 | SciPy 2020 | VanderPlas（提供更多细节）
• NeurIPS 2020：JAX 生态系统聚会（DeepMind 团队讨论围绕 JAX 的生态系统库）
• 面向机器学习及其他领域的 JAX 入门（一次很好的实践型研讨会）
• 第一天演讲：JAX、Flax 和 Transformer | HuggingFace（四场演讲都很精彩）
• 第二天演讲：JAX、Flax 和 Transformer | HuggingFace（只有前两场演讲相关）

博客

• 使用 JAX 加速我们的研究 | DeepMind（与 NeuroIPS 2020 视频内容相似）
• 你并不了解 JAX | Colin Raffel

致谢

• 本笔记本深受官方 JAX、Flax 和 Haiku 文档的启发。

引用

如果你觉得这些内容对你有帮助，请引用如下：

@misc{Gordic2021GetStartedWithJAX,
  author = {Gordić, Aleksa},
  title = {开始使用 JAX},
  year = {2021},
  publisher = {GitHub},
  journal = {GitHub 仓库},
  howpublished = {\url{https://github.com/gordicaleksa/get-started-with-JAX}},
}

与我联系

如果你希望生活中多一些与人工智能相关的内容 :nerd_face:，不妨考虑：

• 订阅我的 YouTube 频道 The AI Epiphany :bell:
• 在 LinkedIn 和 Twitter 上关注我 :bulb:
• 在 Medium 上关注我 :books: :heart:
• 加入 Discord 社区！ :family:

许可证

JAX 快速上手指南

本指南基于 get-started-with-JAX 项目整理，旨在帮助开发者快速入门 JAX、Flax 和 Haiku 生态系统。JAX 正成为 PyTorch 和 TensorFlow 之外日益流行的替代方案，特别适用于高性能机器学习研究。

环境准备

系统要求

• 操作系统：推荐 Linux 或 macOS（Windows 用户建议直接使用 Google Colab 以避免环境配置问题）
• Python 版本：3.8 - 3.11
• 硬件加速：支持 CPU、GPU (CUDA) 和 TPU

前置依赖

确保已安装以下基础工具：

• Python pip 包管理器
• CUDA Toolkit 和 cuDNN（如需 GPU 支持）

安装步骤

1. 安装 JAX (CPU 版本)

pip install jax jaxlib

2. 安装 JAX (GPU 版本)

pip install "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_releases.html

国内加速提示：如果上述谷歌源连接缓慢，可尝试使用清华源安装基础包，但 jaxlib 的 CUDA 版本通常仍需从官方源获取。

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

3. 安装高层库 (可选但推荐)

# 安装 Flax (Google 推出的神经网络库)
pip install flax

# 安装 Haiku (DeepMind 推出的神经网络库)
pip install dm-haiku

# 安装优化器库
pip install optax

4. 免环境配置方案 (推荐 Windows 用户)

直接点击项目中的 Jupyter Notebook 链接，在 Google Colab 中打开运行。这种方式无需本地配置环境，且免费提供 TPU/GPU 资源。

基本使用

以下是最简单的 JAX 代码示例，展示其核心特性：即时编译 (jit)、自动微分 (grad) 和向量化 (vmap)。

示例：纯 JAX 实现简单计算与梯度

import jax
import jax.numpy as jnp

# 1. 定义一个普通函数
def sigmoid(x):
    return 1 / (1 + jnp.exp(-x))

# 2. 使用 jit 进行即时编译加速
fast_sigmoid = jax.jit(sigmoid)

# 3. 使用 grad 自动计算梯度
grad_sigmoid = jax.grad(sigmoid)

# 4. 执行计算
x = jnp.array([0.0, 1.0, -1.0])
y = fast_sigmoid(x)
dy_dx = grad_sigmoid(x)

print(f"输入：{x}")
print(f"Sigmoid 输出：{y}")
print(f"梯度输出：{dy_dx}")

# 5. 使用 vmap 自动向量化 (对批次数据操作)
batch_x = jnp.random.normal(size=(4, 3)) # 4 个样本，每个 3 维
batch_y = jax.vmap(sigmoid)(batch_x)
print(f"批量输出形状：{batch_y.shape}")

下一步学习路径

根据原项目推荐，您可以按顺序深入学习：

1. 从零到英雄 (Zero to Hero)：深入理解 jit, grad, vmap 等核心概念。
2. 多机训练：学习如何在多 TPU/GPU 核心上训练模型。
3. 手写神经网络：不使用框架，仅用纯 JAX 构建并训练 MLP。
4. Flax 实战：学习使用 Flax 构建模块化神经网络。