从 Haiku 迁移到 Flax

内容

从 Haiku 迁移到 Flax#

本指南将展示 Haiku、Flax Linen 和 Flax NNX 之间的区别。Haiku 和 Linen 都强制执行无状态模块的功能范式，而 NNX 是一个新的下一代 API，它拥抱 Python 语言以提供更直观的开发体验。

提升转换#

Flax 和 Haiku 都提供了一组转换，我们将其称为提升转换，它们以这样一种方式包装 JAX 转换：它们可以与模块一起使用，有时还可以提供额外的功能。在本节中，我们将看看如何在 Flax 和 Haiku 中使用 scan 的提升版本来实现一个简单的 RNN 层。

首先，我们将定义一个 RNNCell 模块，它将包含 RNN 单步逻辑。我们还将定义一个 initial_state 方法，该方法将用于初始化 RNN 的状态（也称为 carry）。与 jax.lax.scan 一样，RNNCell.__call__ 方法将是一个接收 carry 和输入并返回新 carry 和输出的函数。在这种情况下，carry 和输出相同。

Haiku

class RNNCell(hk.Module):
  def __init__(self, hidden_size: int, name=None):
    super().__init__(name=name)
    self.hidden_size = hidden_size

  def __call__(self, carry, x):
    x = jnp.concatenate([carry, x], axis=-1)
    x = hk.Linear(self.hidden_size)(x)
    x = jax.nn.relu(x)
    return x, x

  def initial_state(self, batch_size: int):
    return jnp.zeros((batch_size, self.hidden_size))

Linen

class RNNCell(nn.Module):
  hidden_size: int


  @nn.compact
  def __call__(self, carry, x):
    x = jnp.concatenate([carry, x], axis=-1)
    x = nn.Dense(self.hidden_size)(x)
    x = jax.nn.relu(x)
    return x, x

  def initial_state(self, batch_size: int):
    return jnp.zeros((batch_size, self.hidden_size))

NNX

class RNNCell(nnx.Module):
  def __init__(self, input_size, hidden_size, rngs):
    self.linear = nnx.Linear(hidden_size + input_size, hidden_size, rngs=rngs)
    self.hidden_size = hidden_size

  def __call__(self, carry, x):
    x = jnp.concatenate([carry, x], axis=-1)
    x = self.linear(x)
    x = jax.nn.relu(x)
    return x, x

  def initial_state(self, batch_size: int):
    return jnp.zeros((batch_size, self.hidden_size))

接下来，我们将定义一个 RNN 模块，它将包含整个 RNN 的逻辑。在 Haiku 中，我们将首先初始化 RNNCell，然后用它来构造 carry，最后使用 hk.scan 在输入序列上运行 RNNCell。

在 Linen 中，我们将使用 nn.scan 定义一个新的临时类型，它将包装 RNNCell。在此过程中，我们还将指定指示 nn.scan 广播 params 集合（所有步骤共享相同的参数）并且不拆分 params rng 流（以便所有步骤都使用相同的参数初始化），最后我们将指定我们希望 scan 在输入的第二个轴上运行，并将输出沿第二个轴堆叠起来。然后我们将立即使用这个临时类型来创建一个提升的 RNNCell 实例，并用它来创建 carry，并运行 __call__ 方法，它将在序列上进行 scan。

在 NNX 中，我们定义了一个 scan 函数 scan_fn，它将使用在 __init__ 中定义的 RNNCell 来扫描序列。

Haiku

class RNN(hk.Module):
  def __init__(self, hidden_size: int, name=None):
    super().__init__(name=name)
    self.hidden_size = hidden_size

  def __call__(self, x):
    cell = RNNCell(self.hidden_size)
    carry = cell.initial_state(x.shape[0])
    carry, y = hk.scan(
      cell, carry,
      jnp.swapaxes(x, 1, 0)
    )
    y = jnp.swapaxes(y, 0, 1)
    return y

Linen

class RNN(nn.Module):
  hidden_size: int


  @nn.compact
  def __call__(self, x):
    rnn = nn.scan(
      RNNCell, variable_broadcast='params',
      split_rngs={'params': False}, in_axes=1, out_axes=1
    )(self.hidden_size)
    carry = rnn.initial_state(x.shape[0])
    carry, y = rnn(carry, x)

    return y

NNX

class RNN(nnx.Module):
  def __init__(self, input_size: int, hidden_size: int, rngs: nnx.Rngs):
    self.hidden_size = hidden_size
    self.cell = RNNCell(input_size, self.hidden_size, rngs=rngs)

  def __call__(self, x):
    scan_fn = lambda carry, cell, x: cell(carry, x)
    carry = self.cell.initial_state(x.shape[0])
    carry, y = nnx.scan(
      scan_fn, in_axes=(nnx.Carry, None, 1), out_axes=(nnx.Carry, 1)
    )(carry, self.cell, x)

    return y

通常，Flax 和 Haiku 之间的提升转换的主要区别在于，在 Haiku 中，提升转换不操作状态，也就是说，Haiku 将处理 params 和 state，使其在转换内部和外部保持相同的形状。在 Flax 中，提升转换可以同时操作变量集合和 rng 流，用户必须根据转换的语义定义每个转换如何处理不同的集合。

与之前示例一样，参数必须通过 .init() 初始化并传递到 .apply() 以在 Haiku 和 Linen 中进行前向传播。在 NNX 中，参数已被急切地初始化并绑定到有状态模块，并且该模块可以简单地调用输入以进行前向传播。

Haiku

x = jnp.ones((3, 12, 32))

def forward(x):
  return RNN(64)(x)

model = hk.without_apply_rng(hk.transform(forward))

params = model.init(
  random.key(0),
  x=jnp.ones((3, 12, 32)),
)

y = model.apply(
  params,
  x=jnp.ones((3, 12, 32)),
)

Linen

x = jnp.ones((3, 12, 32))




model = RNN(64)

params = model.init(
  random.key(0),
  x=jnp.ones((3, 12, 32)),
)['params']

y = model.apply(
  {'params': params},
  x=jnp.ones((3, 12, 32)),
)

NNX

x = jnp.ones((3, 12, 32))

model = RNN(x.shape[2], 64, rngs=nnx.Rngs(0))

y = model(x)

...

与前几节中的示例相比，唯一值得注意的变化是，这次我们在 Haiku 中使用了 hk.without_apply_rng，因此我们不必将 rng 参数作为 None 传递到 apply 方法。

扫描层#

scan 的一个非常重要的应用是，将一系列层迭代地应用于输入，并将每层的输出作为下一层的输入。这对于减少大型模型的编译时间非常有用。作为示例，我们将创建一个简单的 Block 模块，然后在将 Block 模块应用 num_layers 次的 MLP 模块中使用它。

在 Haiku 中，我们按照通常的方式定义 Block 模块，然后在 MLP 内部，我们将使用 hk.experimental.layer_stack 在一个 stack_block 函数上创建一个 Block 模块堆栈。

在 Linen 中，Block 的定义略有不同，__call__ 将接受并返回一个第二个虚拟输入/输出，在这两种情况下都将是 None。在 MLP 中，我们将使用 nn.scan，如前一个示例中一样，但通过设置 split_rngs={'params': True} 和 variable_axes={'params': 0}，我们告诉 nn.scan 为每个步骤创建不同的参数，并在第一个轴上切片 params 集合，有效地实现了一个 Block 模块堆栈，如 Haiku 中所示。

在 NNX 中，我们使用 nnx.Scan.constructor() 定义一个 Block 模块堆栈。然后我们可以简单地调用 Block 的堆栈 self.blocks，在输入和 carry 上，以获得前向传播输出。

Haiku

class Block(hk.Module):
  def __init__(self, features: int, name=None):
    super().__init__(name=name)
    self.features = features

  def __call__(self, x, training: bool):
    x = hk.Linear(self.features)(x)
    x = hk.dropout(hk.next_rng_key(), 0.5 if training else 0, x)
    x = jax.nn.relu(x)
    return x

class MLP(hk.Module):
  def __init__(self, features: int, num_layers: int, name=None):
      super().__init__(name=name)
      self.features = features
      self.num_layers = num_layers



  def __call__(self, x, training: bool):
    @hk.experimental.layer_stack(self.num_layers)
    def stack_block(x):
      return Block(self.features)(x, training)

    stack = hk.experimental.layer_stack(self.num_layers)
    return stack_block(x)

Linen

class Block(nn.Module):
  features: int
  training: bool

  @nn.compact
  def __call__(self, x, _):
    x = nn.Dense(self.features)(x)
    x = nn.Dropout(0.5)(x, deterministic=not self.training)
    x = jax.nn.relu(x)
    return x, None

class MLP(nn.Module):
  features: int
  num_layers: int




  @nn.compact
  def __call__(self, x, training: bool):
    ScanBlock = nn.scan(
      Block, variable_axes={'params': 0}, split_rngs={'params': True},
      length=self.num_layers)

    y, _ = ScanBlock(self.features, training)(x, None)
    return y

NNX

class Block(nnx.Module):
  def __init__(self, input_dim, features, rngs):
    self.linear = nnx.Linear(input_dim, features, rngs=rngs)
    self.dropout = nnx.Dropout(0.5, rngs=rngs)

  def __call__(self, x: jax.Array, _):
    x = self.linear(x)
    x = self.dropout(x)
    x = jax.nn.relu(x)
    return x, None

class MLP(nnx.Module):
  def __init__(self, input_dim, features, num_layers, rngs):
    self.blocks = nnx.Scan.constructor(
      Block, length=num_layers
    )(input_dim, features, rngs=rngs)



  def __call__(self, x):




    y, _ = self.blocks(x, None)
    return y

请注意，在 Flax 中，我们将 None 作为第二个参数传递给 ScanBlock 并忽略其第二个输出。这些表示每步的输入/输出，但它们是 None，因为在这种情况下我们没有任何输入/输出。

初始化每个模型与之前的示例相同。在这种情况下，我们将指定要使用 5 个层，每个层都有 64 个特征。与之前一样，我们还为 NNX 传入输入形状。

Haiku

def forward(x, training: bool):
  return MLP(64, num_layers=5)(x, training)

model = hk.transform(forward)

sample_x = jnp.ones((1, 64))
params = model.init(
  random.key(0),
  sample_x, training=False # <== inputs
)

Linen

...


model = MLP(64, num_layers=5)

sample_x = jnp.ones((1, 64))
params = model.init(
  random.key(0),
  sample_x, training=False # <== inputs
)['params']

NNX

...

model = MLP(64, 64, num_layers=5, rngs=nnx.Rngs(0))

...

使用层扫描时，您应该注意的一件事是，所有层都被融合成单个层，其参数在第一个轴上有一个额外的“层”维度。在这种情况下，所有参数的形状都将以 (5, ...) 开头，因为我们使用了 5 个层。

Haiku

...


{
    'mlp/__layer_stack_no_per_layer/block/linear': {
        'b': (5, 64),
        'w': (5, 64, 64)
    }
}



...

Linen

...


FrozenDict({
    ScanBlock_0: {
        Dense_0: {
            bias: (5, 64),
            kernel: (5, 64, 64),
        },
    },
})

...

NNX

_, params, _ = nnx.split(model, nnx.Param, ...)

params
State({
  'blocks': {
    'scan_module': {
      'linear': {
        'bias': VariableState(type=Param, value=(5, 64)),
        'kernel': VariableState(type=Param, value=(5, 64, 64))
      }
    }
  }
})

顶层 Haiku 函数与顶层 Flax 模块#

在 Haiku 中，可以使用原始的 hk.{get,set}_{parameter,state} 来定义/访问模型参数和状态，从而将整个模型编写为单个函数。将顶层“模块”编写为函数非常常见。

Flax 团队推荐使用更以模块为中心的方案，该方案使用 __call__ 来定义前向函数。在 Linen 中，相应的访问器将是 Module.param 和 Module.variable（有关集合的解释，请参见处理状态）。在 NNX 中，可以使用常规的 Python 类语义来设置和访问参数和变量。

Haiku

...

def forward(x):

  counter = hk.get_state('counter', shape=[], dtype=jnp.int32, init=jnp.ones)
  multiplier = hk.get_parameter(
    'multiplier', shape=[1,], dtype=x.dtype, init=jnp.ones
  )

  output = x + multiplier * counter

  hk.set_state("counter", counter + 1)
  return output

model = hk.transform_with_state(forward)

params, state = model.init(random.key(0), jnp.ones((1, 64)))

Linen

...


class FooModule(nn.Module):
  @nn.compact
  def __call__(self, x):
    counter = self.variable('counter', 'count', lambda: jnp.ones((), jnp.int32))
    multiplier = self.param(
      'multiplier', nn.initializers.ones_init(), [1,], x.dtype
    )

    output = x + multiplier * counter.value
    if not self.is_initializing():  # otherwise model.init() also increases it
      counter.value += 1
    return output

model = FooModule()
variables = model.init(random.key(0), jnp.ones((1, 64)))
params, counter = variables['params'], variables['counter']

NNX

class Counter(nnx.Variable):
  pass

class FooModule(nnx.Module):

  def __init__(self, rngs):
    self.counter = Counter(jnp.ones((), jnp.int32))
    self.multiplier = nnx.Param(
      nnx.initializers.ones(rngs.params(), [1,], jnp.float32)
    )
  def __call__(self, x):
    output = x + self.multiplier * self.counter.value

    self.counter.value += 1
    return output

model = FooModule(rngs=nnx.Rngs(0))

_, params, counter = nnx.split(model, nnx.Param, Counter)

从 Haiku 迁移到 Flax

内容

从 Haiku 迁移到 Flax#

基本示例#

处理状态#

使用多种方法#

提升转换#

扫描层#

顶层 Haiku 函数与顶层 Flax 模块#