flax.training 包

flax.training 包#

训练状态#

class flax.training.train_state.TrainState(step, apply_fn, params, tx, opt_state)[source]#

适用于使用单个 Optax 优化器的常见情况的简单训练状态。

示例用法

>>> import flax.linen as nn
>>> from flax.training.train_state import TrainState
>>> import jax, jax.numpy as jnp
>>> import optax

>>> x = jnp.ones((1, 2))
>>> y = jnp.ones((1, 2))
>>> model = nn.Dense(2)
>>> variables = model.init(jax.random.key(0), x)
>>> tx = optax.adam(1e-3)

>>> state = TrainState.create(
...     apply_fn=model.apply,
...     params=variables['params'],
...     tx=tx)

>>> def loss_fn(params, x, y):
...   predictions = state.apply_fn({'params': params}, x)
...   loss = optax.l2_loss(predictions=predictions, targets=y).mean()
...   return loss
>>> loss_fn(state.params, x, y)
Array(3.3514676, dtype=float32)

>>> grads = jax.grad(loss_fn)(state.params, x, y)
>>> state = state.apply_gradients(grads=grads)
>>> loss_fn(state.params, x, y)
Array(3.343844, dtype=float32)

请注意，您可以通过子类化轻松扩展此数据类以存储其他数据（例如，其他变量集合）。

对于更奇特的用例（例如多个优化器），最好是派生该类并进行修改。

参数

  • step – 计数器从 0 开始，每次调用 .apply_gradients() 时递增。

  • apply_fn – 通常设置为 model.apply()。 保留在该数据类中是为了方便，以便在训练循环中的 train_step() 函数中拥有更短的参数列表。

  • params – 将由 tx 更新并由 apply_fn 使用的参数。

  • tx – Optax 梯度变换。

  • opt_state – 用于 tx 的状态。

apply_gradients(*, grads, **kwargs)[source]#

更新 step、params、opt_state 和返回值中的 **kwargs。

请注意，在内部，此函数调用 .tx.update()，然后调用 optax.apply_updates() 来更新 params 和 opt_state。

参数

  • grads – 与 .params 具有相同 pytree 结构的梯度。

  • **kwargs – 应 .replace() 的其他数据类属性。

返回

更新后的 self 实例，其中 step 递增 1，params 和 opt_state 通过应用 grads 进行更新，并且其他属性根据 kwargs 进行替换。

classmethod create(*, apply_fn, params, tx, **kwargs)[source]#

使用 step=0 和初始化的 opt_state 创建一个新实例。

内容