Flax 是什么？

Flax 是一个由谷歌开发的高性能、开源的 JAX 神经网络库。它被设计为灵活、可扩展，并与 JAX 的可组合函数转换生态系统（如 jit（即时编译）、grad（自动微分）和 vmap（自动矢量化））无缝协作。与传统的面向对象框架不同，Flax 采用函数式方法，其中模型参数被显式处理，而不是模型对象的内部状态，这使得模型无状态且更易于在 JAX 中管理。

主要特点

基于 JAX 构建： 直接利用 JAX 的速度和强大功能，在 CPU、GPU 和 TPU 上进行高性能计算。
函数式与显式： 模型被定义为无状态的 Python 数据结构（PyTrees），其参数被显式处理。这使得对模型的推理和转换变得简单明了。
灵活性： 提供了一个简单而强大的 API，用于定义复杂的神经网络架构，从简单的 MLP 到大规模的 Transformer。
生态系统集成： 与 JAX 生态系统中的其他库（如用于优化的 Optax 和用于可靠测试的 Chex）良好协作。
经大规模验证： 在谷歌内部用于大规模研究和生产工作负载，确保其健壮性和高性能。

使用案例

深度学习研究： 非常适合需要实现新颖架构并充分利用 JAX 性能和函数式范例的研究人员。
自然语言处理 (NLP)： 构建和训练大型语言模型 (LLM) 和其他基于 Transformer 的架构。
计算机视觉： 开发用于图像分类、目标检测和其他视觉任务的高性能模型。
大规模训练： 在多个加速器（GPU/TPU）上高效地扩展模型训练。

入门指南

这是一个在 Flax 中定义简单多层感知器 (MLP) 的 “Hello World” 示例。

首先，安装 Flax 和 JAX： ```bash pip install flax jax

接下来，定义一个简单的模型并初始化其参数： ```python import jax import jax.numpy as jnp from flax import linen as nn

定义一个简单的 MLP 模型

class SimpleMLP(nn.Module): features: list[int]

@nn.compact
def __call__(self, x):
    for feat in self.features[:-1]:
        x = nn.relu(nn.Dense(features=feat)(x))
    x = nn.Dense(features=self.features[-1])(x)
    return x

创建一个 PRNG 密钥用于初始化

key = jax.random.PRNGKey(0)

定义输入形状

input_shape = (1, 10) # 批量大小为 1，10 个特征

定义模型架构

model = SimpleMLP(features=[64, 32, 1])

初始化模型参数

params = model.init(key, jnp.ones(input_shape))[‘params’]

打印参数的结构

print(jax.tree_util.tree_map(lambda x: x.shape, params))

执行一次前向传播

output = model.apply({‘params’: params}, jnp.ones(input_shape)) print(“输出:”, output)

定价

Flax 是一个完全免费的开源项目，根据 Apache 2.0 许可证发布。

Flax