jax.example_libraries.optimizers 模块#

JAX 优化器编写示例。

您可能并不想导入此模块!此库中的优化器仅供示例使用。如果您正在寻找功能齐全的优化器库,请考虑 Optax

此模块包含一些方便的优化器定义,特别是初始化和更新函数,可用于 ndarray 或任意嵌套的元组/列表/字典 ndarray。

优化器被建模为 (init_fun, update_fun, get_params) 函数三元组,其中组件函数具有以下签名

init_fun(params)

Args:
  params: pytree representing the initial parameters.

Returns:
  A pytree representing the initial optimizer state, which includes the
  initial parameters and may also include auxiliary values like initial
  momentum. The optimizer state pytree structure generally differs from that
  of `params`.
update_fun(step, grads, opt_state)

Args:
  step: integer representing the step index.
  grads: a pytree with the same structure as `get_params(opt_state)`
    representing the gradients to be used in updating the optimizer state.
  opt_state: a pytree representing the optimizer state to be updated.

Returns:
  A pytree with the same structure as the `opt_state` argument representing
  the updated optimizer state.
get_params(opt_state)

Args:
  opt_state: pytree representing an optimizer state.

Returns:
  A pytree representing the parameters extracted from `opt_state`, such that
  the invariant `params == get_params(init_fun(params))` holds true.

请注意,优化器实现对 opt_state 的形式具有很大的灵活性:它只需要是 JaxType 的 pytree(以便可以将其传递给 api.py 中定义的 JAX 变换),并且必须可以被 update_fun 和 get_params 使用。

使用示例

opt_init, opt_update, get_params = optimizers.sgd(learning_rate)
opt_state = opt_init(params)

def step(step, opt_state):
  value, grads = jax.value_and_grad(loss_fn)(get_params(opt_state))
  opt_state = opt_update(step, grads, opt_state)
  return value, opt_state

for i in range(num_steps):
  value, opt_state = step(i, opt_state)
class jax.example_libraries.optimizers.JoinPoint(subtree)[source]#

基类:object

标记两个连接(嵌套)pytree 之间的边界。

class jax.example_libraries.optimizers.Optimizer(init_fn, update_fn, params_fn)[source]#

基类:NamedTuple

参数:
  • init_fn (InitFn)

  • update_fn (UpdateFn)

  • params_fn (ParamsFn)

init_fn: InitFn#

字段编号 0 的别名

params_fn: ParamsFn#

字段编号 2 的别名

update_fn: UpdateFn#

字段编号 1 的别名

class jax.example_libraries.optimizers.OptimizerState(packed_state, tree_def, subtree_defs)#

基类:tuple

packed_state#

字段编号 0 的别名

subtree_defs#

字段编号 2 的别名

tree_def#

字段编号 1 的别名

jax.example_libraries.optimizers.adagrad(step_size, momentum=0.9)[source]#

构建 Adagrad 的优化器三元组。

在线学习和随机优化自适应次梯度方法:http://www.jmlr.org/papers/volume12/duchi11a/duchi11a.pdf

参数:
  • step_size – 正标量,或表示步长计划的可调用对象,该对象将迭代索引映射到正标量。

  • momentum – 可选,动量的正标量值

返回值:

一个 (init_fun, update_fun, get_params) 三元组。

jax.example_libraries.optimizers.adam(step_size, b1=0.9, b2=0.999, eps=1e-08)[source]#

构建 Adam 的优化器三元组。

参数:
  • step_size – 正标量,或表示步长计划的可调用对象,该对象将迭代索引映射到正标量。

  • b1 – 可选,beta_1 的正标量值,第一矩估计的指数衰减率(默认为 0.9)。

  • b2 – 可选,beta_2 的正标量值,第二矩估计的指数衰减率(默认为 0.999)。

  • eps – 可选,epsilon 的正标量值,一个小常数,用于数值稳定性(默认为 1e-8)。

返回值:

一个 (init_fun, update_fun, get_params) 三元组。

jax.example_libraries.optimizers.adamax(step_size, b1=0.9, b2=0.999, eps=1e-08)[source]#

构建 AdaMax 的优化器三元组(基于无限范数的 Adam 变体)。

参数:
  • step_size – 正标量,或表示步长计划的可调用对象,该对象将迭代索引映射到正标量。

  • b1 – 可选,beta_1 的正标量值,第一矩估计的指数衰减率(默认为 0.9)。

  • b2 – 可选,beta_2 的正标量值,第二矩估计的指数衰减率(默认为 0.999)。

  • eps – 可选,epsilon 的正标量值,一个小常数,用于数值稳定性(默认为 1e-8)。

返回值:

一个 (init_fun, update_fun, get_params) 三元组。

jax.example_libraries.optimizers.clip_grads(grad_tree, max_norm)[source]#

将存储为数组 pytree 的梯度裁剪到最大范数 max_norm

jax.example_libraries.optimizers.constant(step_size)[source]#
返回类型:

计划

jax.example_libraries.optimizers.exponential_decay(step_size, decay_steps, decay_rate)[source]#
jax.example_libraries.optimizers.inverse_time_decay(step_size, decay_steps, decay_rate, staircase=False)[source]#
jax.example_libraries.optimizers.l2_norm(tree)[source]#

计算数组 pytree 的 l2 范数。对权重衰减很有用。

jax.example_libraries.optimizers.make_schedule(scalar_or_schedule)[source]#
参数:

scalar_or_schedule (float | Schedule)

返回类型:

计划

jax.example_libraries.optimizers.momentum(step_size, mass)[source]#

构建带动量的 SGD 的优化器三元组。

参数:
  • step_size (Schedule) – 正标量,或表示步长计划的可调用对象,该对象将迭代索引映射到正标量。

  • mass (float) – 表示动量系数的正标量。

返回值:

一个 (init_fun, update_fun, get_params) 三元组。

jax.example_libraries.optimizers.nesterov(step_size, mass)[source]#

构建带 Nesterov 动量的 SGD 的优化器三元组。

参数:
  • step_size (Schedule) – 正标量,或表示步长计划的可调用对象,该对象将迭代索引映射到正标量。

  • mass (float) – 表示动量系数的正标量。

返回值:

一个 (init_fun, update_fun, get_params) 三元组。

jax.example_libraries.optimizers.optimizer(opt_maker)[source]#

装饰器,使为数组定义的优化器能够推广到容器。

使用此装饰器,您可以编写 init、update 和 get_params 函数,这些函数分别仅对单个数组进行操作,并将它们转换为相应的功能,这些功能对参数的 pytree 进行操作。有关示例,请参阅 optimizers.py 中定义的优化器。

参数:

opt_maker (Callable[..., tuple[Callable[[Params], State], Callable[[Step, Updates, Params], Params], Callable[[State], Params]]]) –

一个返回 (init_fun, update_fun, get_params) 函数三元组的函数,这些函数可能仅适用于 ndarrays,如

init_fun :: ndarray -> OptStatePytree ndarray
update_fun :: OptStatePytree ndarray -> OptStatePytree ndarray
get_params :: OptStatePytree ndarray -> ndarray

返回值:

一个 (init_fun, update_fun, get_params) 函数三元组,适用于任意 pytrees,如

init_fun :: ParameterPytree ndarray -> OptimizerState
update_fun :: OptimizerState -> OptimizerState
get_params :: OptimizerState -> ParameterPytree ndarray

返回函数使用的 OptimizerState pytree 类型与 ParameterPytree (OptStatePytree ndarray) 同构,但为了性能,可能会将状态存储为例如部分扁平化的数据结构。

返回类型:

Callable[…, Optimizer]

jax.example_libraries.optimizers.pack_optimizer_state(marked_pytree)[source]#

将标记的 pytree 转换为 OptimizerState。

unpack_optimizer_state 的逆操作。将外部 pytree 的叶子表示为 JoinPoint 的标记 pytree 转换回 OptimizerState。此函数旨在用于反序列化优化器状态。

参数:

marked_pytree – 包含保存更多 pytree 的 JoinPoint 叶子的 pytree。

返回值:

与输入参数等效的 OptimizerState。

jax.example_libraries.optimizers.piecewise_constant(boundaries, values)[source]#
参数:
  • boundaries (Any)

  • values (Any)

jax.example_libraries.optimizers.polynomial_decay(step_size, decay_steps, final_step_size, power=1.0)[source]#
jax.example_libraries.optimizers.rmsprop(step_size, gamma=0.9, eps=1e-08)[source]#

构建 RMSProp 的优化器三元组。

参数:

step_size – 正标量,或表示步长计划的可调用对象,该计划将迭代索引映射到正标量。gamma:衰减参数。eps:Epsilon 参数。

返回值:

一个 (init_fun, update_fun, get_params) 三元组。

jax.example_libraries.optimizers.rmsprop_momentum(step_size, gamma=0.9, eps=1e-08, momentum=0.9)[source]#

构建带动量的 RMSProp 的优化器三元组。

此优化器与 rmsprop 优化器分开,因为它需要跟踪其他参数。

参数:
  • step_size – 正标量,或表示步长计划的可调用对象,该对象将迭代索引映射到正标量。

  • gamma – 衰减参数。

  • eps – Epsilon 参数。

  • momentum – 动量参数。

返回值:

一个 (init_fun, update_fun, get_params) 三元组。

jax.example_libraries.optimizers.sgd(step_size)[source]#

构建随机梯度下降的优化器三元组。

参数:

step_size – 正标量,或表示步长计划的可调用对象,该对象将迭代索引映射到正标量。

返回值:

一个 (init_fun, update_fun, get_params) 三元组。

jax.example_libraries.optimizers.sm3(step_size, momentum=0.9)[source]#

构建 SM3 的优化器三元组。

用于大规模学习的内存高效自适应优化。 https://arxiv.org/abs/1901.11150

参数:
  • step_size – 正标量,或表示步长计划的可调用对象,该对象将迭代索引映射到正标量。

  • momentum – 可选,动量的正标量值

返回值:

一个 (init_fun, update_fun, get_params) 三元组。

jax.example_libraries.optimizers.unpack_optimizer_state(opt_state)[source]#

将 OptimizerState 转换为标记的 pytree。

将 OptimizerState 转换为标记的 pytree,其中外部 pytree 的叶子表示为 JoinPoint 以避免丢失信息。此函数旨在用于序列化优化器状态。

参数:

opt_state – OptimizerState

返回值:

一个 pytree,其 JoinPoint 叶子包含第二级 pytree。