Function adam_optimizer

axlearn/common/optimizers.py:859–878

Adam optimizer with l2 regularization.

(
    learning_rate: schedule.Schedule,
    *,
    b1: float,
    b2: float,
    eps: float,
    l2_regularizer_weight: float = 0,
    l2_regularizer_per_param_scale: Optional[Callable[[NestedOptParam], Any]] = None,
    mu_dtype: Optional[jnp.dtype] = None,
)

def adam_optimizer(
    learning_rate: schedule.Schedule,
    *,
    b1: float,
    b2: float,
    eps: float,
    l2_regularizer_weight: float = 0,
    l2_regularizer_per_param_scale: Optional[Callable[[NestedOptParam], Any]] = None,
    mu_dtype: Optional[jnp.dtype] = None,
) -> PartitionedGradientTransformation:
    """Adam optimizer with l2 regularization."""
    tx = [
        l2_regularizer(
            regularizer_weight=l2_regularizer_weight,
            per_param_scale=l2_regularizer_per_param_scale,
        ),
        adam_partition(optax.scale_by_adam(b1=b1, b2=b2, eps=eps, mu_dtype=mu_dtype)),
        scale_by_schedule(scale_from_learning_rate(learning_rate)),
    ]
    return chain(*tx)


class EmaState(NamedTuple):