Flax NNX 与 JAX 变换

Flax NNX 与 JAX 变换#

注意力

本页面与新的 Flax NNX API 相关。

在本指南中，您将了解使用 Flax NNX 和 JAX 变换的差异，以及如何在它们之间无缝切换或将它们一起使用。在本指南中，我们将重点介绍 jit 和 grad 函数变换。

首先，让我们设置导入并生成一些虚拟数据

from flax import nnx
import jax

x = jax.random.normal(jax.random.key(0), (1, 2))
y = jax.random.normal(jax.random.key(1), (1, 3))

NNX 与 JAX 变换之间的差异#

Flax NNX 与 JAX 变换之间的主要区别在于，Flax NNX 变换允许您变换接受 Flax NNX 图对象作为参数（Module、Rngs、Optimizer 等）的函数，即使这些函数的状态将被修改，而它们在 JAX 变换中不被识别。因此，Flax NNX 变换可以变换不纯净的函数，并进行修改和副作用。

Flax NNX 的 函数式 API 提供了一种将图结构转换为 pytrees 并返回的方法。通过在每个函数边界执行此操作，您可以有效地将图结构与任何 JAX 变换一起使用，并以与函数纯度一致的方式传播状态更新。Flax NNX 自定义变换，例如 nnx.jit 和 nnx.grad，只是删除了样板代码，因此代码看起来是带状态的。

以下是使用 nnx.jit 和 nnx.grad 变换与使用 jax.jit 和 jax.grad 变换的示例。请注意，Flax NNX 变换函数的函数签名可以直接接受 nnx.Linear 模块，并且可以对该模块进行有状态的更新，而 JAX 变换函数的函数签名只能接受 pytree 注册的 State 和 GraphDef 对象，并且必须返回它们的更新副本以维护变换函数的纯度。

@nnx.jit
def train_step(model, x, y):
  def loss_fn(model):
    return ((model(x) - y) ** 2).mean()
  grads = nnx.grad(loss_fn)(model)
  params = nnx.state(model, nnx.Param)
  params = jax.tree_util.tree_map(
    lambda p, g: p - 0.1 * g, params, grads
  )
  nnx.update(model, params)

model = nnx.Linear(2, 3, rngs=nnx.Rngs(0))
train_step(model, x, y)

@jax.jit
def train_step(graphdef, state, x, y):
  def loss_fn(graphdef, state):
    model = nnx.merge(graphdef, state)
    return ((model(x) - y) ** 2).mean()
  grads = jax.grad(loss_fn, argnums=1)(graphdef, state)

  model = nnx.merge(graphdef, state)
  params = nnx.state(model, nnx.Param)
  params = jax.tree_util.tree_map(
    lambda p, g: p - 0.1 * g, params, grads
  )
  nnx.update(model, params)
  return nnx.split(model)

graphdef, state = nnx.split(nnx.Linear(2, 3, rngs=nnx.Rngs(0)))
graphdef, state = train_step(graphdef, state, x, y)

混合使用 Flax NNX 和 JAX 变换#

只要 JAX 变换函数是纯净的，并且具有 JAX 识别的有效参数类型，就可以将 Flax NNX 和 JAX 变换混合在一起。

@nnx.jit
def train_step(model, x, y):
  def loss_fn(graphdef, state):
    model = nnx.merge(graphdef, state)
    return ((model(x) - y) ** 2).mean()
  grads = jax.grad(loss_fn, 1)(*nnx.split(model))
  params = nnx.state(model, nnx.Param)
  params = jax.tree_util.tree_map(
    lambda p, g: p - 0.1 * g, params, grads
  )
  nnx.update(model, params)

model = nnx.Linear(2, 3, rngs=nnx.Rngs(0))
train_step(model, x, y)

@jax.jit
def train_step(graphdef, state, x, y):
  model = nnx.merge(graphdef, state)
  def loss_fn(model):
    return ((model(x) - y) ** 2).mean()
  grads = nnx.grad(loss_fn)(model)
  params = nnx.state(model, nnx.Param)
  params = jax.tree_util.tree_map(
    lambda p, g: p - 0.1 * g, params, grads
  )
  nnx.update(model, params)
  return nnx.split(model)

graphdef, state = nnx.split(nnx.Linear(2, 3, rngs=nnx.Rngs(0)))
graphdef, state = train_step(graphdef, state, x, y)

目录