Flax：专为灵活性设计的 JAX 神经网络库及生态系统

概述 | 快速安装 | Flax 是什么样子的？ | 文档

2024 年发布的 Flax NNX 是一个简化的全新 Flax API，旨在让在 JAX 中创建、检查、调试和分析神经网络变得更加容易。它通过为 Python 引用语义提供一等公民级的支持来实现这一点。这使得用户能够使用常规 Python 对象来表达他们的模型，从而支持引用共享和可变性。

Flax NNX 源自 Flax Linen API。该 API 于 2020 年由 Google Brain 的工程师和研究人员与 JAX 团队紧密合作发布。

您可以在 专门的 Flax 文档网站 上了解更多关于 Flax NNX 的信息。请务必查看：

• Flax NNX 基础
• MNIST 教程
• 为什么选择 Flax NNX
• 从 Flax Linen 到 Flax NNX 的演进

注意： Flax Linen 的 文档有其独立的站点。

Flax 团队的使命是服务于不断增长的 JAX 神经网络研究生态系统——无论是在 Alphabet 内部还是更广泛的社区中——并探索 JAX 发挥优势的用例。我们几乎所有的协调与规划工作都在 GitHub 上进行，同时也在那里讨论即将发生的设计变更。我们欢迎对任何讨论、问题或拉取请求线程提出反馈意见。

您可以在我们的 Flax GitHub 讨论区 提出功能请求、告知我们您的工作进展、报告问题或提问。

我们预计会持续改进 Flax，但不打算对核心 API 进行重大破坏性更改。我们会尽可能使用 Changelog 条目和弃用警告。

如果您想直接联系我们，我们的邮箱是 flax-dev@google.com。

概述

Flax 是一个面向 JAX 的高性能神经网络库及生态系统，其设计宗旨就是 灵活性： 尝试新的训练方式，只需分叉示例并修改训练循环，而无需向框架添加新特性。

Flax 在与 JAX 团队的密切合作下开发，自带您开展研究所需的一切，包括：

• 神经网络 API (flax.nnx)：包含 Linear、Conv、BatchNorm、LayerNorm、GroupNorm、Attention（如 MultiHeadAttention）、LSTMCell、GRUCell、Dropout。

• 实用工具与模式：分布式训练、序列化与检查点保存、指标计算、设备端预取。

• 教学示例：MNIST、使用 Gemma 语言模型（Transformer）进行推理/采样。

快速安装

Flax 使用 JAX，因此请先查看 JAX 在 CPU、GPU 和 TPU 上的安装说明。

您需要 Python 3.8 或更高版本。从 PyPI 安装 Flax：

pip install flax

要升级到最新版本的 Flax，可以使用：

pip install --upgrade git+https://github.com/google/flax.git

如果某些依赖项需要额外的包（例如 matplotlib），而这些包并未包含在内，您可以运行以下命令来安装所有依赖项：

pip install "flax[all]"

Flax 是什么样子的？

我们提供了三个使用 Flax API 的示例：一个简单的多层感知机、一个卷积神经网络和一个自编码器。

要深入了解 Module 抽象，可以查看我们的文档以及我们关于 Module 抽象的详细介绍。如需更多最佳实践的具体演示，请参阅我们的 指南和 开发者笔记。

多层感知机的示例：

class MLP(nnx.Module):
  def __init__(self, din: int, dmid: int, dout: int, *, rngs: nnx.Rngs):
    self.linear1 = nnx.Linear(din, dmid, rngs=rngs)
    self.dropout = nnx.Dropout(rate=0.1, rngs=rngs)
    self.bn = nnx.BatchNorm(dmid, rngs=rngs)
    self.linear2 = nnx.Linear(dmid, dout, rngs=rngs)

  def __call__(self, x: jax.Array):
    x = nnx.gelu(self.dropout(self.bn(self.linear1(x))))
    return self.linear2(x)

卷积神经网络的示例：

class CNN(nnx.Module):
  def __init__(self, *, rngs: nnx.Rngs):
    self.conv1 = nnx.Conv(1, 32, kernel_size=(3, 3), rngs=rngs)
    self.conv2 = nnx.Conv(32, 64, kernel_size=(3, 3), rngs=rngs)
    self.avg_pool = partial(nnx.avg_pool, window_shape=(2, 2), strides=(2, 2))
    self.linear1 = nnx.Linear(3136, 256, rngs=rngs)
    self.linear2 = nnx.Linear(256, 10, rngs=rngs)

  def __call__(self, x):
    x = self.avg_pool(nnx.relu(self.conv1(x)))
    x = self.avg_pool(nnx.relu(self.conv2(x)))
    x = x.reshape(x.shape[0], -1)  # 展平
    x = nnx.relu(self.linear1(x))
    x = self.linear2(x)
    return x

自编码器的示例：

Encoder = lambda rngs: nnx.Linear(2, 10, rngs=rngs)
Decoder = lambda rngs: nnx.Linear(10, 2, rngs=rngs)

class AutoEncoder(nnx.Module):
  def __init__(self, rngs):
    self.encoder = Encoder(rngs)
    self.decoder = Decoder(rngs)

  def __call__(self, x) -> jax.Array:
    return self.decoder(self.encoder(x))

  def encode(self, x) -> jax.Array:
    return self.encoder(x)

引用 Flax

引用本仓库时，请使用以下格式：

@software{flax2020github,
  author = {Jonathan Heek and Anselm Levskaya and Avital Oliver and Marvin Ritter and Bertrand Rondepierre and Andreas Steiner and Marc van {Z}ee},
  title = {{F}lax: A neural network library and ecosystem for {JAX}},
  url = {http://github.com/google/flax},
  version = {0.12.6},
  year = {2024},
}

在上述 BibTeX 条目中，作者姓名按字母顺序排列；版本号应为 flax/version.py 中的版本号；年份对应于该项目的开源发布年份。

注意事项

Flax 是由 Google DeepMind 的一支专业团队维护的开源项目，但它并非 Google 的官方产品。

Flax 快速上手指南

Flax 是一个为 JAX 设计的高性能神经网络库，主打灵活性。2024 年推出的 Flax NNX 是新简化版 API，支持原生 Python 引用语义，让模型定义、调试和分析更加直观。

环境准备

在开始之前，请确保满足以下要求：

• 操作系统：Linux, macOS, 或 Windows (WSL2 推荐)
• Python 版本：3.8 或更高
• 前置依赖：必须先安装 JAX。
  • 请访问 JAX 安装指南 根据您的硬件（CPU/GPU/TPU）安装对应版本的 JAX。
  • 国内用户提示：如果官方源下载缓慢，可尝试使用清华或阿里镜像源安装 JAX 相关包。

安装步骤

1. 基础安装

通过 PyPI 安装稳定版：

pip install flax

2. 安装最新版（可选）

如需体验最新功能，可从 GitHub 源码安装：

pip install --upgrade git+https://github.com/google/flax.git

3. 安装完整依赖（推荐）

若需使用绘图等额外功能（如 matplotlib），建议安装完整包：

pip install "flax[all]"

国内加速提示：若 pip 下载速度慢，可添加国内镜像源参数，例如：

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

基本使用

Flax NNX 允许你像编写普通 Python 类一样定义神经网络。以下是一个最简单的多层感知机（MLP）示例：

import jax
import jax.numpy as jnp
from flax import nnx

class MLP(nnx.Module):
  def __init__(self, din: int, dmid: int, dout: int, *, rngs: nnx.Rngs):
    # 定义网络层
    self.linear1 = nnx.Linear(din, dmid, rngs=rngs)
    self.dropout = nnx.Dropout(rate=0.1, rngs=rngs)
    self.bn = nnx.BatchNorm(dmid, rngs=rngs)
    self.linear2 = nnx.Linear(dmid, dout, rngs=rngs)

  def __call__(self, x: jax.Array):
    # 定义前向传播逻辑
    x = nnx.gelu(self.dropout(self.bn(self.linear1(x))))
    return self.linear2(x)

# 实例化模型
rngs = nnx.Rngs(0)
model = MLP(din=10, dmid=20, dout=2, rngs=rngs)

# 创建虚拟输入数据
x = jnp.ones((4, 10))

# 执行前向推理
output = model(x)
print(output.shape)  # 输出: (4, 2)

核心要点：

• 继承 nnx.Module 类来定义模型。
• 在 __init__ 中初始化层，需传入 rngs 用于随机数生成。
• 在 __call__ 中定义数据流动逻辑。
• 直接调用模型实例即可进行计算。