in jcm/losses.py [0:0]
def step_fn(carry_state, batch):
"""Running one step of training or evaluation.
This function will undergo `jax.lax.scan` so that multiple steps can be pmapped and jit-compiled together
for faster execution.
Args:
carry_state: A tuple (JAX random state, `flax.struct.dataclass` containing the training state).
batch: A mini-batch of training/evaluation data.
Returns:
new_carry_state: The updated tuple of `carry_state`.
loss: The average loss value of this state.
"""
(rng, state) = carry_state
rng, step_rng = jax.random.split(rng)
grad_fn = jax.value_and_grad(loss_fn, argnums=1, has_aux=True)
if train:
step = state.step
params = state.params
states = state.model_state
opt_state = state.opt_state
target_ema, num_scales = ema_scales_fn(step)
if target_ema is None and num_scales is None:
(
loss,
(new_model_state, log_stats),
), grad = grad_fn(step_rng, params, states, batch)
grad = jax.lax.pmean(grad, axis_name="batch")
new_params, new_opt_state = optimize_fn(grad, opt_state, params)
new_params_ema = jax.tree_util.tree_map(
lambda p_ema, p: p_ema * state.ema_rate
+ p * (1.0 - state.ema_rate),
state.params_ema,
new_params,
)
step = state.step + 1
new_state = state.replace(
step=step,
params=new_params,
params_ema=new_params_ema,
model_state=new_model_state,
opt_state=new_opt_state,
)
else:
target_params = state.target_params
(loss, (new_model_state, log_stats)), grad = grad_fn(
step_rng, params, states, batch, target_params, num_scales
)
grad = jax.lax.pmean(grad, axis_name="batch")
new_params, new_opt_state = optimize_fn(grad, opt_state, params)
new_params_ema = jax.tree_util.tree_map(
lambda p_ema, p: p_ema * state.ema_rate
+ p * (1.0 - state.ema_rate),
state.params_ema,
new_params,
)
new_target_params = jax.tree_util.tree_map(
lambda p_target, p: p_target * target_ema + p * (1.0 - target_ema),
target_params,
new_params,
)
step = state.step + 1
new_state = state.replace(
step=step,
params=new_params,
params_ema=new_params_ema,
target_params=new_target_params,
model_state=new_model_state,
opt_state=new_opt_state,
)
else:
target_ema, num_scales = ema_scales_fn(state.step)
if target_ema is None and num_scales is None:
loss, (_, log_stats) = loss_fn(
step_rng,
state.params_ema,
state.model_state,
batch,
)
else:
loss, (_, log_stats) = loss_fn(
step_rng,
state.params_ema,
state.model_state,
batch,
state.target_params,
num_scales,
)
new_state = state
loss = jax.lax.pmean(loss, axis_name="batch")
mean_log_stats = jax.tree_map(
lambda x: jax.lax.pmean(x, axis_name="batch"), log_stats
)
new_carry_state = (rng, new_state)
return new_carry_state, (loss, mean_log_stats)