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
的形式具有很大的灵活性:它只需要是 JaxTypes 的 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
标记两个连接(嵌套)pytrees 之间的边界。
- 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.inverse_time_decay(step_size, decay_steps, decay_rate, staircase=False)[source]#
- 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]#
用于使针对数组定义的优化器泛化到容器的装饰器。
使用此装饰器,您可以编写仅作用于单个数组的初始化、更新和获取参数函数,并将其转换为作用于参数 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)
函数三元组,可作用于任意 pytree,根据以下定义: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 的叶子表示为 JoinPoints 的标记 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) 三元组。