Simple train state for the common case with a single Optax optimizer.
>>> 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(1.8136346, dtype=float32)
>>> grads = jax.grad(loss_fn)(state.params, x, y)
>>> state = state.apply_gradients(grads=grads)
>>> loss_fn(state.params, x, y)
Array(1.8079796, dtype=float32)
Note that you can easily extend this dataclass by subclassing it for storing additional data (e.g. additional variable collections).
For more exotic usecases (e.g. multiple optimizers) it’s probably best to fork the class and modify it.
RetroSearch is an open source project built by @garambo | Open a GitHub Issue
Search and Browse the WWW like it's 1997 | Search results from DuckDuckGo
HTML:
3.2
| Encoding:
UTF-8
| Version:
0.7.4