flax.training 包

训练状态

class flax.training.train_state.TrainState(step, apply_fn, params, tx, opt_state)

对于具有单个 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)

更新返回值中的 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)

创建一个新的实例，其中 step=0 并初始化 opt_state。