timm documentation

优化

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

优化

此页面包含 timm 中优化器的 API 参考文档,并概述了如何通过优化器工厂选择和配置这些优化器。

用法

优化器通过名称使用 create_optimizer_v2 创建。给定一个模型,工厂会先构建参数组(包括权重衰减排除项和可选的逐层学习率衰减),然后再构建优化器。优化器特定的参数会作为关键字参数传递。

import timm
from timm.optim import create_optimizer_v2, list_optimizers

model = timm.create_model('vit_base_patch16_224')

# AdamW. 1D params (biases, norm scales) and params in model.no_weight_decay() are excluded from weight decay.
optimizer = create_optimizer_v2(model, opt='adamw', lr=1e-3, weight_decay=0.05)

# Layer-wise learning rate decay (fine-tuning), per-layer lr scales from the model's group_matcher().
optimizer = create_optimizer_v2(model, opt='adamw', lr=1e-3, weight_decay=0.05, layer_decay=0.75)

# Optimizer specific arguments are passed through.
optimizer = create_optimizer_v2(model, opt='lamb', lr=5e-3, weight_decay=0.02, trust_clip=True)

# A list of params or param groups can be passed instead of a model, weight decay grouping is then up to the caller.
optimizer = create_optimizer_v2(model.parameters(), opt='sgd', lr=0.1, momentum=0.9)

# All registered names, optionally with a description.
print(list_optimizers(with_description=True))

训练脚本(train.py)通过以下标志暴露了相同的功能:

  • --opt 用于选择优化器。
  • --lr 用于设置学习率。
  • --weight-decay 用于设置权重衰减。
  • --momentum 用于设置动量。
  • --opt-eps 用于设置 epsilon(eps)。
  • --opt-betas 用于设置 beta 系数(betas)。
  • --layer-decay 设置逐层学习率衰减。
  • --opt-kwargs key=value ... 用于传递其他优化器参数。

命名约定

已注册名称为小写,由基础优化器名称加上修饰符组成:

修改含义示例
c 前缀谨慎变体,当更新量的符号与梯度不一致时,该更新会被屏蔽(Cautious Optimizers)cadamw, clion, cmars, cadopt
w 后缀解耦(AdamW 风格)权重衰减,适用于默认使用耦合 L2 衰减的优化器sgdw, lambw, adoptw, novogradw, rmsproptfw
c 后缀修正后的权重衰减,衰减系数按 lr² / max_lr 而非 lr 进行缩放(为什么梯度在训练末期会迅速增大)adamc, nadamc, sgdc, lionc, adoptc, kronc
n 前缀Nesterov 动量(SGD、LARS 和 Muon 系列)nesterov, nlars, nmuon, nadamuon
lookahead_ 前缀将优化器包裹在 Lookahead 中lookahead_adamw
fused 前缀NVIDIA Apex 融合实现(需要 apex)fusedadamw, fusedlamb
bnb 前缀bitsandbytes 实现,8bit 后缀表示 8 位状态(需要 bitsandbytes)bnbadamw, bnbadamw8bit

有几个名称早于这些约定:lambc 和 larc 是带有信任比裁剪的 LAMB / LARS(修正衰减的 LAMB 变体为 lambcd),momentum 是带有经典动量的 SGD,而 sgd 是带有 Nesterov 动量的 SGD。下表中 legacy 的实现早于它们被纳入 PyTorch 的时间。

优化器系列

名称实现备注
adamw, adam, nadam, radam, radamw, adamax, adadelta, adagrad, rmsprop, sgd, momentum, nesterovtorch.optim
adamwlegacy, adamc, cadamw, cadamcAdamWLegacy谨慎使用 AdamW,修正衰减,更新裁剪,capturable
nadamw, nadamc, cnadamw, cnadamcNAdamW带有 Nesterov 动量的 AdamW(MLCommons 算法效率形式)
adopt, adoptw, adoptc, cadopt, cadoptw, cadoptcAdopt适用于任意 β2 的 Adam 变体
adan, adanw, cadan, cadanwAdan自适应 Nesterov 动量
adabelief, radabeliefAdaBeliefradabelief 添加了 RAdam 风格的校正
adafactor, cadafactorAdafactorfairseq Adafactor,当 lr=None 时的相对步长
adafactorbv, adafactorbvc, cadafactorbv, cadafactorbvcAdafactorBigVision大型视觉 Adafactor,分解的二阶矩,低精度动量,更新裁剪
adamp, cadamp, sgdp, csgdpAdamP, SGDP尺度不变权重的更新投影
adahessianAdahessian二阶,需要一个闭包(is_second_order)
kron, kronw, kroncKron使用 Kronecker 分解预条件器的 PSGD
lamb, lambc, lambw, lambcw, lambcd, clamb, clambwLamb逐层自适应矩,trust_clip 变体(lambc、lambcw)
laprop, lapropc, clapropLaProp动量与自适应性分离
lars, larc, nlars, nlarcLars逐层自适应速率缩放,trust_clip 变体(larc、nlarc)
lion, lionc, clion, clioncLion演化符号动量
madgrad, madgradwMADGRAD基于动量的双重平均
mars, cmarsMars方差缩减,AdamW 或 Lion 风格的更新
muon, nmuon, adamuon, nadamuonMuon针对 2D 及以上权重使用 Newton-Schulz 正交化动量,针对 1D 参数回退到 AdamW / NAdamW
novograd, novogradw, novogradcNvNovoGrad逐层归一化矩
rmsproptf, rmsproptfw, rmsproptfc, crmsproptfRMSpropTFTensorFlow 中算子的排序和 eps 的处理方式,学习率被折叠到动量缓冲区中
sgdw, nesterovw, sgdc, csgdw, csgdcSGDW带有解耦权重衰减的 SGD
nadamlegacy, radamlegacyNAdamLegacy, RAdamLegacy参考实现,已被torch.optim取代
fused*, bnb*NVIDIA Apex,bitsandbytes在软件包安装时注册

常用选项

大多数 timm 优化器除了常见的 lr、weight_decay、betas 和 eps 之外,还共享一组选项。具体哪些优化器支持哪些选项,可以查看下方的类签名;工厂函数会对不支持的参数抛出异常。

  • caution 应用谨慎更新掩码,等效于带有 c 前缀的注册表名称。
  • corrected_weight_decay 使用 lr² / max_lr 而非 lr(即带 c 后缀的名称)来缩放解耦权重衰减,从而使有效衰减率跟随学习率调度。 max_lr 是构建优化器时使用的学习率,请将其设置为调度中的峰值学习率(如 train.py 所示)。该值保存在每个参数组中(max_lr_snapshot),并随优化器状态一起保存和恢复。
  • decoupled_decay(或 decoupled)将默认使用耦合 L2 衰减的优化器切换为 AdamW 风格的衰减(即带有 w 后缀的名称)。
  • foreach 选择多张量实现。默认情况下(None),在 CUDA 及其他具有 foreach 内核的设备上使用该实现,与 torch.optim 保持一致,除非配置不支持(例如没有 capturable 的张量 lr,或在缺少所需 _foreach_maximum_ 重载的 PyTorch 版本上使用 caution)。多张量步骤按参数分块运行,每个块的大小经过调整,使其工作集保持在 GPU L2 缓存中,对于大型模型,这比每次操作都将所有参数通过内存流式传输一次更快。块大小由设备 L2 缓存大小推导得出(PyTorch >= 2.5),并且在结果块低于 4M 个元素的设备上会跳过分块。TIMM_FOREACH_CHUNK_SIZE=<elements> 用于覆盖块大小,0 用于禁用分块。
  • capturable(AdamWLegacy、NAdamW、Adopt、Muon)将步数计数保留在设备上,以便在 CUDA 图中捕获该步。所有优化器都接受以 0 维张量形式给出的 lr,该张量由调度器就地更新,以便与 torch.compile 和 CUDA 图配合使用。
  • clipping_threshold(AdamWLegacy、NAdamW、AdafactorBigVision)通过其 RMS 对更新量进行裁剪,公式为 update / max(1, rms / threshold),其中 rms 为二阶矩归一化梯度的 RMS,按张量计算(StableAdamW、Adafactor)。取值与论文中一致,为 1.0。与 Adafactor 相同,(解耦的)权重衰减不会被裁剪,而 StableAdamW 还会使用裁剪后的学习率对衰减进行缩放。
  • maximize 用于最大化目标函数,而非最小化。

低精度参数(bfloat16、float16)在整个过程中均受支持:

  • AdaBelief、RAdamLegacy 和 Adafactor 会为低精度参数在 float32 中保留其状态。
  • AdafactorBigVision 会在降低的精度(momentum_dtype)下保存其动量。
  • Kron 具有独立的 mu_dtype / precond_dtype 选项。

这些状态 dtype 会被 load_state_dict 保留(PyTorch 的默认加载器会将状态转换为参数 dtype)。 step 函数避免了主机同步(对设备张量进行数据依赖的 Python 分支),AdamP / SGDP 的投影检查是例外。

参数组

给定一个模型,create_optimizer_v2 会使用 param_groups_weight_decay 对参数进行分组(1D 参数以及模型 no_weight_decay() 集合中的参数不应用衰减),或者在设置了 layer_decay 时,使用 param_groups_layer_decay,后者利用模型的 group_matcher() 为每一层分配学习率缩放因子 layer_decay ** (num_layers - layer_id)。没有 group_matcher() 的模型会回退到 auto_group_layers,这需要模型的 pretrained_cfg['classifier'] 来分离头部。对于混合优化器(Muon),fallback_list 和 fallback_no_weight_decay 参数会将参数路由到回退(AdamW)优化器。

API

工厂函数

timm.optim.create_optimizer_v2

< >

( model_or_params: typing.Union[torch.nn.Module, torch.optim.optimizer.ParamsT]opt: str = 'sgd'lr: typing.Optional[float] = Noneweight_decay: float = 0.0momentum: float = 0.9foreach: typing.Optional[bool] = Nonefilter_bias_and_bn: bool = Truefallback_list: typing.Collection[str] = ()fallback_no_weight_decay: bool = Falselayer_decay: typing.Optional[float] = Nonelayer_decay_min_scale: float = 0.0layer_decay_no_opt_scale: typing.Optional[float] = Noneparam_group_fn: typing.Optional[typing.Callable[[torch.nn.Module], torch.optim.optimizer.ParamsT]] = None**kwargs: typing.Any )

参数

  • model_or_params — PyTorch 模型或参数/参数组的可迭代对象。 如果提供了模型,将根据其他参数自动提取并分组参数。
  • opt — 要创建的优化器名称(例如 ‘adam’、‘adamw’、‘sgd’)。 使用 list_optimizers() 查看可用选项。
  • lr — 学习率。如果为 None,将使用优化器的默认值。
  • weight_decay — 权重衰减系数。如果 model_or_params 是一个模型,将用于创建参数组。
  • momentum — 支持动量的优化器的动量因子。仅当所选优化器接受动量参数时使用。
  • foreach — 如果可用,启用/禁用 foreach(多张量)实现。 如果为 None,将使用特定于优化器的默认值。
  • filter_bias_and_bn — 如果为 True,则不会对 bias 和归一化层参数(所有 1d 参数)应用权重衰减。仅当 model_or_params 是模型且 weight_decay > 0 时使用。
  • fallback_list — 用于混合优化器(例如 Muon 的 AdamW)的参数名称模式集合,以使用回退优化器。支持通配符匹配。
  • fallback_no_weight_decay — 如果为 True,模型 no_weight_decay() 列表中的参数将使用混合优化器的回退优化器(例如 Muon 的 AdamW)。
  • layer_decay — 可选的逐层学习率衰减因子。如果提供,学习率将按 layer_decay^(max_depth - layer_depth) 进行缩放。仅在 model_or_params 为模型时使用。
  • param_group_fn — 用于创建自定义参数组的可选函数。 如果提供,其他参数分组选项将被忽略。
  • **kwargs — 其他优化器特定的参数(例如 Adam 的 betas)。

通过 timm 注册表创建优化器实例。

创建并配置具有适当参数组和设置的优化器。 支持为权重衰减和逐层学习率自动创建参数组,也支持自定义参数分组。

示例:

从模型中获取 AdamW 时,参数组会自动创建:

optimizer = create_optimizer_v2(model, 'adamw', lr=1e-3)

带有动量和权重衰减的 SGD:

optimizer = create_optimizer_v2(model, 'sgd', lr=0.1, momentum=0.9, weight_decay=1e-4)

采用逐层学习率衰减的 Adam:

optimizer = create_optimizer_v2(model, 'adam', lr=1e-3, layer_decay=0.7)

自定义参数组:

def group_fn(model):
    return [
        {'params': model.backbone.parameters(), 'lr': 1e-4},
        {'params': model.head.parameters(), 'lr': 1e-3},
    ]

optimizer = create_optimizer_v2(model, 'sgd', param_group_fn=group_fn)

注意: 参数组处理优先级:

  1. 如果提供了 param_group_fn,则仅会使用它
  2. 如果提供了 layer_decay,则会创建逐层分组
  3. 如果 weight_decay > 0 且 filter_bias_and_bn 为 True,则会创建权重衰减组
  4. 否则,所有参数将位于同一个组中

timm.optim.list_optimizers

< >

( filter: typing.Union[str, typing.List[str]] = ''exclude_filters: typing.Optional[typing.List[str]] = Nonewith_description: bool = False ) → If with_description is False

参数

  • filter — 通配符风格的过滤字符串或过滤字符串列表 (例如,‘adam’ 表示所有 Adam 变体,或 [‘adam’, ‘*8bit’] 表示 Adam 变体和 8 位优化器)。空字符串表示不进行过滤。
  • exclude_filters — 可选的通配符模式列表,用于排除特定项。例如, [’8bit’, ‘fused’] 将排除 8 位和融合实现。
  • with_description — 如果为 True,则返回 (name, description) 元组,而不仅仅是名称。描述提供了对优化器特征的简要说明。

返回值

如果 with_description 为 False

优化器名称的字符串列表(例如,[‘adam’, ‘adamw’, …]) 如果 with_description 为 True: (name, description) 元组的列表(例如,[(‘adam’, ‘Adaptive Moment…’), …])

列出可用的优化器名称,可选择进行过滤。

列出所有已注册的优化器,支持使用通配符模式进行可选过滤。 优化器可以使用 include 和 exclude 模式进行过滤,并且可以选择在返回每个优化器名称的同时返回其描述。

示例:

>>> list_optimizers()
['adam', 'adamw', 'sgd', ...]

>>> list_optimizers(['la*', 'nla*'])  # List lamb & lars
['lamb', 'lambc', 'larc', 'lars', 'nlarc', 'nlars']

>>> list_optimizers('*adam*', exclude_filters=['bnb*', 'fused*'])  # Exclude bnb & apex adam optimizers
['adam', 'adamax', 'adamp', 'adamw', 'nadam', 'nadamw', 'radam']

>>> list_optimizers(with_description=True)  # Get descriptions
[('adabelief', 'Adapts learning rate based on gradient prediction error'),
 ('adadelta', 'torch.optim Adadelta, Adapts learning rates based on running windows of gradients'),
 ('adafactor', 'Memory-efficient implementation of Adam with factored gradients'),
...]

timm.optim.get_optimizer_class

< >

( name: strbind_defaults: bool = True ) → If bind_defaults is False

参数

  • name — 要获取的优化器名称(例如,‘adam’、‘sgd’)
  • bind_defaults — 如果为 True,则返回一个绑定了来自 OptimInfo 的默认参数的偏函数。 如果为 False,则返回原始的优化器类。

返回值

如果 bind_defaults 为 False

优化器类(例如 torch.optim.Adam) 如果 bind_defaults 为 True: 一个绑定了默认参数的偏函数

引发

ValueError

  • ValueError — 如果在注册表中找不到优化器名称

按名称获取优化器类,并可选择绑定默认参数。

获取优化器类或绑定了默认参数的偏函数。 这允许直接以默认配置实例化优化器,而无需通过完整的工厂。

示例:

绑定注册表默认值后,sgd 已启用 nesterov 动量:

SGD = get_optimizer_class('sgd')
opt = SGD(model.parameters(), lr=0.1, momentum=0.9)

原始优化器类:

SGD = get_optimizer_class('sgd', bind_defaults=False)
opt = SGD(model.parameters(), lr=1e-3, momentum=0.9, nesterov=True)

timm.optim.get_optimizer_info

< >

( name: str )

参数

  • name — 优化器的名称

引发

ValueError

  • ValueError — 如果未找到优化器

获取优化器的 OptimInfo。

参数组函数

timm.optim.param_groups_weight_decay

< >

( model: Moduleweight_decay: typing.Optional[float] = 1e-05weight_decay_exclude_1d: bool = Trueno_weight_decay_list: typing.Collection[str] = ()fallback_list: typing.Collection[str] = ()fallback_no_weight_decay: bool = False )

timm.optim.param_groups_layer_decay

< >

( model: Moduleweight_decay: float = 0.05no_weight_decay_list: typing.Collection[str] = ()fallback_list: typing.Collection[str] = ()fallback_no_weight_decay: bool = Falseweight_decay_exclude_1d: bool = Truelayer_decay: float = 0.75min_scale: float = 0.0no_opt_scale: typing.Optional[float] = Noneverbose: bool = False )

用于逐层学习率衰减和权重衰减的参数组 基于 BEiT:https://github.com/microsoft/unilm/blob/master/beit/optim_factory.py#L58

timm.optim.auto_group_layers

< >

( modellayers_per_group = 12num_groups = None )

优化器类

类 timm.optim.AdaBelief

< >

( paramslr = 0.001betas = (0.9, 0.999)eps = 1e-16weight_decay = 0amsgrad = Falsedecoupled_decay = Truefixed_decay = Falserectify = Truedegenerated_to_sgd = Trueforeach: typing.Optional[bool] = None )

参数

  • params(可迭代对象)— 要优化的参数的可迭代对象,或定义参数组的字典
  • lr (float, 可选) — 学习率(默认值:1e-3)
  • betas (Tuple[float, float], optional) — 用于计算梯度及其平方移动平均的系数(默认值:(0.9, 0.999))
  • eps (float, 可选) — 添加到分母中以提高数值稳定性的项(默认值:1e-16)
  • weight_decay (float, 可选) — 权重衰减(L2 惩罚)(默认值:0)
  • amsgrad(布尔值,可选)— 是否使用来自论文 On the Convergence of Adam and Beyond_ 的该算法的 AMSGrad 变体(默认值:False)
  • decoupled_decay(布尔值,可选)—(默认值:True)如果设置为 True,则优化器将使用 AdamW 中的解耦权重衰减
  • fixed_decay(布尔值,可选)—(默认值:False)当 weightdecouple 设置为 True 时,此参数生效。 当 fixed_decay == True 时,权重衰减的计算方式为 $W{new} = W{old} - W{old} \times decay$。 当 fixeddecay == False 时,权重衰减的计算方式为 $W{new} = W{old} - W{old} \times decay \times lr$。请注意,在这种情况下, 权重衰减比例会随着学习率(lr)的减小而降低。
  • rectify(布尔值,可选)—(默认值:True)如果设置为 True,则执行类似于 RAdam 的 rectified 更新
  • degenerated_to_sgd(布尔值,可选)(默认值 —True)如果设置为 True,则在梯度方差较高时执行 SGD 更新
  • foreach(布尔值,可选)— 是否使用 foreach(多张量)实现,默认(None)情况下,除非不支持,否则会使用

实现了 AdaBelief 算法。基于 PyTorch 中的 Adam 修改而来

reference: AdaBelief 优化器,通过观察到的梯度的置信度来调整步长,NeurIPS 2020

有关推荐超参数的完整表格,请参阅 https://github.com/juntang-zhuang/Adabelief-Optimizer’ 例如,EfficientNet 的 train/args 可参考这些 gists

步骤

< >

( closure = None )

参数

  • closure(可调用对象,可选)— 一个闭包,用于重新评估模型并返回损失。

执行单个优化步骤。

类 timm.optim.Adafactor

< >

( params: ParamsTlr: typing.Optional[float] = Noneeps: float = 1e-30eps_scale: float = 0.001clip_threshold: float = 1.0decay_rate: float = -0.8betas: typing.Optional[typing.Tuple[float, float]] = Noneweight_decay: float = 0.0scale_parameter: bool = Truewarmup_init: bool = Falsemin_dim_size_to_factor: int = 16caution: bool = False )

实现 Adafactor 算法。

此实现基于:Adafactor: Adaptive Learning Rates with Sublinear Memory Cost(参见 https://arxiv.org/abs/1804.04235)

请注意,该优化器会根据 scale_parameter、relative_step 和 warmup_init 选项在内部调整学习率。

当设置了 lr 时,会使用手动(外部)学习率调度(此时 relative_step 为 False)。 请注意,在此实现中,scale_parameter 和 warmup_init 仅适用于相对步长学习率(lr=None)。与 fairseq / HF Transformers 不同,参数 RMS 缩放不会应用于外部学习率。

Ags: params: 要优化的参数可迭代对象,或定义参数组的字典 lr: 外部学习率 eps: 分别用于平方梯度和参数尺度的正则化常数 eps_scale: 分别用于参数尺度的正则化常数 clip_threshold: 最终梯度更新的均方根阈值 decay_rate: 用于计算平方梯度移动平均的系数 beta1: 用于计算梯度移动平均的系数 weight_decay: 权重衰减 scale_parameter: 如果为 True,相对步长学习率将根据参数的均方根进行缩放 warmup_init: 时间依赖的学习率计算取决于是否使用预热初始化

步骤

< >

( closure = None )

参数

  • closure(可调用对象,可选)— 一个闭包,用于重新评估模型并返回损失。

执行单个优化步骤。

class timm.optim.AdafactorBigVision

< >

( params: ParamsTlr: float = 1.0min_dim_size_to_factor: int = 16decay_rate: float = 0.8decay_offset: int = 0beta2_cap: float = 0.999momentum: typing.Optional[float] = 0.9momentum_dtype: typing.Union[str, torch.dtype] = torch.bfloat16eps: typing.Optional[float] = Noneweight_decay: float = 0.0clipping_threshold: typing.Optional[float] = Noneunscaled_wd: bool = Falsecaution: bool = Falsecorrected_weight_decay: bool = False )

PyTorch 实现的 BigVision Adafactor 变体,同时提供单张量和多张量实现。

改编自 Ross Wightman 的 https://github.com/google-research/big_vision

类 timm.optim.Adahessian

< >

( paramslr = 0.1betas = (0.9, 0.999)eps = 1e-08weight_decay = 0.0hessian_power = 1.0update_each = 1n_samples = 1avg_conv_kernel = False )

参数

  • params(可迭代对象)— 要优化的参数可迭代对象,或定义参数组的字典
  • lr(float,可选)— 学习率(默认值:0.1)
  • betas ((float, float), optional) — 用于计算梯度和平方海森迹的滑动平均的系数(默认值:(0.9, 0.999))
  • eps (float, 可选) — 添加到分母中以提高数值稳定性的项(默认值:1e-8)
  • weight_decay (float, 可选) — 权重衰减(L2 惩罚)(默认值:0.0)
  • hessian_power(float,可选)— 海森矩阵迹的指数(默认值:1.0)
  • update_each (int, 可选) — 仅在经过 此 数量的步骤后计算海森矩阵迹的近似值(以节省时间)(默认值:1)
  • n_samples (int, 可选) — 为近似 hessian 迹而对 z 进行采样的次数(默认值:1)

Implements the AdaHessian algorithm from “ADAHESSIAN: An Adaptive Second OrderOptimizer for Machine Learning”

get_params

< >

( )

获取所有 param_groups 中具有梯度的所有参数

set_hessian

< >

( )

计算海森矩阵迹的 Hutchinson 近似值,并为每个可训练参数进行累加。

步骤

< >

( closure = None )

参数

  • closure(可调用对象,可选)— 一个用于重新评估模型并返回损失的闭包(默认值 — None)

执行单个优化步骤。

zero_hessian

< >

( )

将累积的 hessian 迹清零。

类 timm.optim.AdamP

< >

( paramslr = 0.001betas = (0.9, 0.999)eps = 1e-08weight_decay = 0delta = 0.1wd_ratio = 0.1nesterov = Falsecaution = False )

类 timm.optim.AdamWLegacy

< >

( params: ParamsTlr: float = 0.001betas: typing.Tuple[float, float] = (0.9, 0.999)eps: float = 1e-08weight_decay: float = 0.01amsgrad: bool = Falsecaution: bool = Falsecorrected_weight_decay: bool = Falsemaximize: bool = Falseforeach: typing.Optional[bool] = Nonecapturable: bool = Falseclipping_threshold: typing.Optional[float] = None )

参数

  • params — 要优化的参数迭代器,或定义参数组的字典
  • lr — 学习率
  • betas — 用于计算梯度及其平方的移动平均的系数
  • eps — 添加到分母中以提高数值稳定性的项
  • weight_decay — 权重衰减系数
  • amsgrad — 是否使用论文 On the Convergence of Adam and Beyond 中该算法的 AMSGrad 变体
  • caution — 使用 AdamW 时请谨慎
  • corrected_weight_decay — 应用修正后的权重衰减(lr**2 / max_lr)
  • maximize — 根据目标函数最大化参数,而不是最小化
  • foreach — whether foreach implementation of optimizer is used. If unspecified by the user (so foreach is None), we will try to use foreach over for-loop implementation on CUDA, since it is faster in general.
  • capturable — 此实例是否可安全地捕获到 CUDA 图中。 传入 True 可能会降低未使用图时的性能,因此如果您不打算 通过图捕获此实例,请保持为 False
  • clipping_threshold — RMS 更新裁剪阈值(StableAdamW,https://arxiv.org/abs/2304.13013),更新量会乘以 1 / max(1, rms / clipping_threshold),其中 rms 为 grad / (sqrt(v_hat) + eps) 的均方根值。 当梯度相对于其二阶矩估计出现突增时,会抑制相应参数的更新量。 论文中取值为 1.0,默认值 None 表示禁用。与 Adafactor 类似,解耦的权重衰减不会被裁剪(StableAdamW 还会将衰减量乘以裁剪后的学习率)。

实现 AdamW 算法。

NOTE: 此实现已被弃用,推荐使用 torch.optim.AdamW,此处仅作为参考

参考文献:

步骤

< >

( closure = None )

参数

  • closure(可调用对象,可选)— 一个闭包,用于重新评估模型并返回损失。

执行单个优化步骤。

类 timm.optim.Adan

< >

( paramslr: float = 0.001betas: typing.Tuple[float, float, float] = (0.98, 0.92, 0.99)eps: float = 1e-08weight_decay: float = 0.0no_prox: bool = Falsecaution: bool = Falseforeach: typing.Optional[bool] = None )

参数

  • params — 要优化的参数可迭代对象,或定义参数组的字典。
  • lr — 学习率。
  • betas — 用于一阶和二阶矩的系数。
  • eps — 添加到分母中以提高数值稳定性的项。
  • weight_decay — 解耦权重衰减(L2 惩罚)
  • no_prox — 如何执行权重衰减
  • caution — 从“Cautious Optimizers”中启用 caution
  • foreach — 如果为 True,将使用 torch._foreach 实现。速度更快,但会占用略多的内存。

实现 Adan 的 pytorch 变体。

Adan 算法在论文 Adan: Adaptive Nesterov Momentum Algorithm for Faster Optimizing Deep Models https://arxiv.org/abs/2208.06677 中提出。

步骤

< >

( closure = None )

执行单个优化步骤。

类 timm.optim.Adopt

< >

( params: ParamsTlr: typing.Union[float, torch.Tensor] = 0.001betas: typing.Tuple[float, float] = (0.9, 0.9999)eps: float = 1e-06clip_exp: typing.Optional[float] = 0.333weight_decay: float = 0.0decoupled: bool = Falsecorrected_weight_decay: bool = Falsecaution: bool = Falseforeach: typing.Optional[bool] = Nonemaximize: bool = Falsecapturable: bool = Falsedifferentiable: bool = False )

ADOPT:修改后的 Adam 可以以任意 β2 达到最优收敛速率:https://arxiv.org/abs/2411.02853

步骤

< >

( closure = None )

参数

  • closure (Callable, optional) — 一个闭包,用于重新评估模型并返回损失。

执行一次优化步骤。

class timm.optim.kron.Kron

< >

( params: ParamsTlr: float = 0.001momentum: float = 0.9weight_decay: float = 0.0preconditioner_update_probability: typing.Union[typing.Callable, float, NoneType] = Nonemax_size_triangular: int = 2048min_ndim_triangular: int = 2memory_save_mode: typing.Optional[str] = Nonemomentum_into_precond_update: bool = Trueprecond_lr: float = 0.1precond_init_scale: float = 1.0mu_dtype: typing.Optional[torch.dtype] = Noneprecond_dtype: typing.Optional[torch.dtype] = Nonedecoupled_decay: bool = Falsecorrected_weight_decay: bool = Falseflatten: bool = Falseflatten_start_dim: int = 2flatten_end_dim: int = -1stochastic_weight_decay: bool = Falsedeterministic: bool = False )

参数

  • params — 要优化的参数可迭代对象,或定义参数组的字典。
  • lr — 学习率。
  • momentum — 动量参数。
  • weight_decay — 权重衰减。
  • preconditioner_update_probability — 更新预条件器的概率。 如果为 None,则默认使用一个从 1.0 退火到 0.03 的调度,持续 4000 步。
  • max_size_triangular — 维度预条件器为三角矩阵时的最大尺寸。
  • min_ndim_triangular — 层需要具有三角预条件器的最小维度数。
  • memory_save_mode — ‘one_diag’、‘smart_one_diag’ 或 ‘all_diag’,默认值为 None 将所有预条件器设置为三角矩阵时,‘one_diag’ 会将每层中最大或最后一个维度设置为对角矩阵,而 ‘all_diag’ 会将所有预条件器设置为对角矩阵。
  • momentum_into_precond_update — 是否将动量发送到预条件器更新中,而不是原始梯度。
  • mu_dtype — 动量累加器的数据类型。
  • precond_dtype — 预条件器的数据类型。
  • decoupled_decay — AdamW 风格的解耦权重衰减
  • corrected_weight_decay — 在使用 decoupled_decay 时应用修正后的权重衰减(lr**2 / max_lr)
  • flatten — 展平维度,而不是完全依赖表达式来处理高秩参数
  • flatten_start_dim — 展平范围的起始维度,默认为 2。对于 ConvNets 来说,这是一个不错的折中方案。
  • flatten_end_dim — 展平范围的结束位置,默认为 -1。
  • stochastic_weight_decay — 启用权重衰减的随机调制
  • deterministic — 在保存/加载(恢复)过程中保持确定性行为。FIXME 速度慢,需要优化

实现了来自 https://github.com/lixilinx/psgd_torch 的 PSGD Kron。

类 timm.optim.Lamb

< >

( params: ParamsTlr: float = 0.001bias_correction: bool = Truebetas: typing.Tuple[float, float] = (0.9, 0.999)eps: float = 1e-06weight_decay: float = 0.01grad_averaging: bool = Truemax_grad_norm: typing.Optional[float] = 1.0trust_clip: bool = Falsealways_adapt: bool = Falsecaution: bool = Falsedecoupled_decay: bool = Falsecorrected_weight_decay: bool = Falseforeach: typing.Optional[bool] = None )

参数

  • params — 要优化的参数可迭代对象,或定义参数组的字典。
  • lr — 学习率
  • betas — 用于计算梯度及其范数的移动平均的系数。
  • eps — 添加到分母中以提高数值稳定性的项。
  • weight_decay — 权重衰减
  • grad_averaging — 在计算梯度的移动平均时,是否对梯度应用 (1-beta2)。
  • max_grad_norm — 用于裁剪全局梯度范数的值。
  • trust_clip — 启用 LAMBC 信任比率裁剪。
  • always_adapt — 对 0.0 的 weight decay 参数应用自适应学习率。
  • caution — 请谨慎操作。
  • decoupled — 应用解耦权重衰减
  • corrected_weight_decay — 在使用 decoupled_decay 时,应用修正后的权重衰减(lr**2 / max_lr)
  • foreach — 使用 foreach(多张量)实现,默认情况下(None)会使用该实现,除非不支持

实现了来自 apex.optimizers.FusedLAMB 的 FuseLAMB(NvLamb 变体)优化器的纯 PyTorch 版本 参考:https://github.com/NVIDIA/DeepLearningExamples/blob/master/PyTorch/LanguageModeling/Transformer-XL/pytorch/lamb.py

LAMB 由以下文献提出:

步骤

< >

( closure = None )

参数

  • closure(可调用对象,可选)— 一个闭包,用于重新评估模型并返回损失。

执行单个优化步骤。

类 timm.optim.LaProp

< >

( params: ParamsTlr: float = 0.0004betas: typing.Tuple[float, float] = (0.9, 0.999)eps: float = 1e-15weight_decay: float = 0.0caution: bool = Falsecorrected_weight_decay: bool = Falseforeach: typing.Optional[bool] = None )

参数

  • foreach — 使用 foreach(多张量)实现,默认情况下(None)会使用该实现,除非不支持(参见 laprop())

LaProp 优化器

论文:LaProp: Separating Momentum and Adaptivity in Adam,https://arxiv.org/abs/2002.04839

步骤

< >

( closure = None )

参数

  • closure(可调用对象,可选)— 一个闭包,用于重新评估模型并返回损失。

执行单个优化步骤。

类 timm.optim.Lars

< >

( paramslr = 1.0momentum = 0dampening = 0weight_decay = 0nesterov = Falsetrust_coeff = 0.001eps = 1e-08trust_clip = Falsealways_adapt = Falseforeach: typing.Optional[bool] = None )

参数

  • params(可迭代对象)— 要优化的参数可迭代对象,或定义参数组的字典。
  • lr (float, 可选) — 学习率(默认值:1.0)。
  • momentum (float, 可选) — 动量因子(默认值:0)
  • weight_decay (float, 可选) — 权重衰减(L2 惩罚)(默认值:0)
  • dampening(float,可选)— 动量的阻尼系数(默认值:0)
  • nesterov(布尔值,可选)— 启用 Nesterov 动量(默认值:False)
  • trust_coeff (float) — 用于计算自适应学习率 / trust_ratio 的信任系数(默认值:0.001)
  • eps (float) — 除法分母的 eps(默认值:1e-8)
  • trust_clip (bool) — 启用 LARC 信任比率裁剪(默认值:False)
  • always_adapt (bool) — 始终应用 LARS 学习率自适应,否则仅在 group weight_decay != 0 时应用(默认值:False)
  • foreach (bool, optional) — 是否使用 foreach(多张量)实现,默认(None)情况下,除非不支持,否则会使用

用于 PyTorch 的 LARS

论文:Large batch training of Convolutional Networks - https://arxiv.org/pdf/1708.03888.pdf

步骤

< >

( closure = None )

参数

  • closure(可调用对象,可选)— 一个闭包,用于重新评估模型并返回损失。

执行单个优化步骤。

类 timm.optim.Lion

< >

( params: ParamsTlr: float = 0.0001betas: typing.Tuple[float, float] = (0.9, 0.99)weight_decay: float = 0.0caution: bool = Falsecorrected_weight_decay: bool = Falsemaximize: bool = Falseforeach: typing.Optional[bool] = None )

实现 Lion 算法。

步骤

< >

( closure = None )

参数

  • closure — 一个闭包,用于重新评估模型并返回损失。

执行单个优化步骤。

类 timm.optim.Lookahead

< >

( base_optimizeralpha = 0.5k = 6 )

类 timm.optim.MADGRAD

< >

( params: typing.Anylr: float = 0.01momentum: float = 0.9weight_decay: float = 0eps: float = 1e-06decoupled_decay: bool = False )

参数

  • params (iterable) — 要优化的参数可迭代对象,或定义参数组的字典。
  • lr (float) — 学习率(默认值:1e-2)。
  • momentum (float) — 取值范围为 [0,1) 的动量值(默认值:0.9)。
  • weight_decay (float) — 权重衰减,即 L2 惩罚(默认值:0)。
  • eps (float) — 添加到根号外分母中的项,用于提高数值稳定性。(默认值:1e-6)。

MADGRAD_:一种用于随机优化的动量化、自适应、对偶平均梯度方法。

.. _MADGRAD: https://arxiv.org/abs/2101.11075

MADGRAD 是一种通用优化器,可替代 SGD 或 Adam,可能收敛更快且泛化能力更强。目前仅支持 GPU。通常,可以使用与 SGD 或 Adam 相同的学习率调度策略。其整体学习率与这两种方法均不可直接比较,应通过超参数搜索来确定。

MADGRAD 所需的权重衰减比其他方法更少,有时甚至可以低至零。用于 SGD 或 Adam 的 beta1 的动量值在此处也应该适用。

对于稀疏问题,weight_decay 和 momentum 都应设置为 0。

步骤

< >

( closure: typing.Optional[typing.Callable[[], float]] = None )

参数

  • closure(可调用对象,可选)— 一个闭包,用于重新评估模型并返回损失。

执行单个优化步骤。

类 timm.optim.Mars

< >

( params: ParamsTlr: float = 0.003betas: typing.Tuple[float, float] = (0.9, 0.99)eps: float = 1e-08weight_decay: float = 0.0gamma: float = 0.025mars_type: str = 'adamw'optimize_1d: bool = Falselr_1d_factor: float = 1.0betas_1d: typing.Optional[typing.Tuple[float, float]] = Nonecaution: bool = False )

MARS 优化器

论文:MARS:释放方差缩减在训练大模型中的潜力 https://arxiv.org/abs/2411.10438

步骤

< >

( closure = None )

参数

  • closure(可调用对象,可选)— 一个闭包,用于重新评估模型并返回损失。

执行单个优化步骤。

类 timm.optim.Muon

< >

( params: ParamsTlr: float = 0.02weight_decay: float = 0momentum: float = 0.95nesterov: bool = Falsens_steps: int = 5ns_coefficients: typing.Union[str, typing.Tuple[float, float, float], typing.List[typing.Tuple[float, float, float]]] = 'quintic'eps: float = 1e-07safety_factor: float = 1.0adjust_lr_fn: typing.Optional[str] = 'match_rms_adamw'conv_mode: str = 'flatten'normalize_spatial: bool = Trueadamw_lr: typing.Optional[float] = Nonefallback_lr_scale: float = 1.0betas: typing.Tuple[float, float] = (0.9, 0.95)algo: str = 'muon'scale_eps: bool = Falsecapturable: bool = Falseverbose: bool = False )

Muon - 通过 Newton-Schulz 正交化的动量

将 Muon 用于 2D 及以上参数(权重矩阵),将 AdamW 用于 1D 参数(偏置、归一化层)以及设置了 ‘use_fallback=True’(或为兼容性设置 ‘use_muon=False’)的参数组。

支持两种算法:

  • “muon”:带有动量 + 正交化的标准 Muon 算法
  • “adamuon”:AdaMuon 算法,在正交化方向上添加逐元素二阶矩估计,以实现类似 Adam 的自适应缩放

步骤

< >

( closure = None )

执行单个优化步骤。

类 timm.optim.NAdamLegacy

< >

( paramslr = 0.002betas = (0.9, 0.999)eps = 1e-08weight_decay = 0schedule_decay = 0.004 )

参数

  • params(可迭代对象)— 要优化的参数的可迭代对象,或定义参数组的字典
  • lr (float, 可选) — 学习率(默认值:2e-3)
  • betas (Tuple[float, float], optional) — 用于计算梯度及其平方的移动平均的系数
  • eps (float, 可选) — 添加到分母中以提高数值稳定性的项(默认值:1e-8)
  • weight_decay (float, 可选) — 权重衰减(L2 惩罚)(默认值:0)
  • schedule_decay(float,可选)— 动量调度衰减(默认值:4e-3)

实现 Nadam 算法(基于 Nesterov 动量的 Adam 变体)。

NOTE: 此实现已被弃用,推荐使用 torch.optim.NAdam,此处仅作为参考保留

Incorporating Nesterov Momentum into Adam__ 中提出了该方案。

http://cs229.stanford.edu/proj2015/054_report.pdf http://www.cs.toronto.edu/~fritz/absps/momentum.pdf

原始出处:https://github.com/pytorch/pytorch/pull/1408 注意:可能存在潜在问题,但在某些问题上表现良好。

步骤

< >

( closure = None )

参数

  • closure(可调用对象,可选)— 一个闭包,用于重新评估模型并返回损失。

执行单个优化步骤。

类 timm.optim.NAdamW

< >

( params: ParamsTlr: float = 0.001betas: typing.Tuple[float, float] = (0.9, 0.999)eps: float = 1e-08weight_decay: float = 0.01caution: bool = Falsecorrected_weight_decay: bool = Falsemaximize: bool = Falseforeach: typing.Optional[bool] = Nonecapturable: bool = Falseclipping_threshold: typing.Optional[float] = None )

参数

  • params — 要优化的参数迭代器,或定义参数组的字典
  • lr — 学习率
  • betas — 用于计算梯度及其平方的移动平均的系数
  • eps — 添加到分母中以提高数值稳定性的项
  • weight_decay — 权重衰减系数
  • caution — 启用 caution
  • corrected_weight_decay — 应用修正后的权重衰减(lr**2 / max_lr)
  • maximize — 根据目标函数最大化参数,而不是最小化
  • foreach — 是否使用优化器的 foreach 实现。如果未指定(None),则在 CUDA(具有 foreach 内核的设备)上使用 foreach 实现,因为它通常更快。
  • capturable — 此实例是否可安全地捕获到 CUDA 图中
  • clipping_threshold — RMS 更新裁剪阈值(StableAdamW,https://arxiv.org/abs/2304.13013),更新量会乘以 1 / max(1, rms / clipping_threshold),其中 rms 为 grad / (sqrt(v_hat) + eps) 的均方根值。 当梯度相对于其二阶矩估计出现突增时,会抑制相应参数的更新量。 论文中取值为 1.0,默认值 None 表示禁用。与 Adafactor 类似,解耦的权重衰减不会被裁剪(StableAdamW 还会将衰减量乘以裁剪后的学习率)。

实现 NAdamW 算法。

有关 NAdam 算法的实现,请参阅 https://arxiv.org/abs/1910.05446 中的表 1(代码中也有注释,指出了 NAdamW 和 AdamW 之间唯一的区别)。

有关该算法的更多详细信息,请参阅

步骤

< >

( closure = None )

参数

  • closure(可调用对象,可选)— 一个闭包,用于重新评估模型并返回损失。

执行单个优化步骤。

类 timm.optim.NvNovoGrad

< >

( paramslr = 0.001betas = (0.95, 0.98)eps = 1e-08weight_decay = 0grad_averaging = Falseamsgrad = Falsedecoupled_decay = Falsecorrected_weight_decay = False )

参数

  • params(可迭代对象)— 要优化的参数的可迭代对象,或定义参数组的字典
  • lr (float, 可选) — 学习率(默认值:1e-3)
  • betas (Tuple[float, float], optional) — 用于计算梯度及其平方的移动平均的系数(默认值:(0.95, 0.98))
  • eps (float, 可选) — 添加到分母中以提高数值稳定性的项(默认值:1e-8)
  • weight_decay (float, 可选) — 权重衰减(L2 惩罚)(默认值:0)
  • grad_averaging — 梯度平均
  • amsgrad(布尔值,可选)— 是否使用来自论文 On the Convergence of Adam and Beyond_ 的该算法的 AMSGrad 变体(默认值:False)
  • decoupled_decay — 将权重衰减直接应用于权重(每步 lr * weight_decay),而不是在动量更新前将其加到归一化梯度上(如论文所述)
  • corrected_weight_decay — 当 decoupled_decay 为 True 时,应用修正后的权重衰减(lr**2 / max_lr)

实现 Novograd 算法。

步骤

< >

( closure = None )

参数

  • closure(可调用对象,可选)— 一个用于重新评估模型的闭包
  • and 返回损失。 —

执行单个优化步骤。

类 timm.optim.RAdamLegacy

< >

( paramslr = 0.001betas = (0.9, 0.999)eps = 1e-08weight_decay = 0 )

PyTorch RAdam 优化器

NOTE: 此实现已被弃用,推荐使用 torch.optim.AdamW,此处仅作为参考

类 timm.optim.RMSpropTF

< >

( params: ParamsTlr: float = 0.01alpha: float = 0.9eps: float = 1e-10weight_decay: float = 0momentum: float = 0.0centered: bool = Falsedecoupled_decay: bool = Falsecorrected_weight_decay: bool = Falselr_in_momentum: bool = Truecaution: bool = Falseforeach: typing.Optional[bool] = None )

参数

  • params — 要优化的参数迭代器,或定义参数组的字典
  • lr — 学习率
  • momentum — 动量因子
  • alpha — 平滑(衰减)常数
  • eps — 添加到分母中以提高数值稳定性的项
  • centered — 如果为 True,则计算居中的 RMSProp,梯度会由其方差的估计值进行归一化
  • weight_decay — 权重衰减(L2 惩罚)(默认值:0)
  • decoupled_decay — 按照 https://arxiv.org/abs/1711.05101 实现的解耦权重衰减
  • corrected_weight_decay — 当 decoupled_decay 为 True 时,应用修正后的权重衰减(lr**2 / max_lr)
  • lr_in_momentum — 学习率缩放包含在动量缓冲区的更新中,遵循 Tensorflow 的默认设置
  • caution — 谨慎操作
  • foreach — 使用 foreach(多张量)实现,默认情况下(None)会使用该实现,除非不支持(参见 rmsprop_tf())

实现 RMSprop 算法(TensorFlow 风格的 epsilon)

NOTE: 这是 PyTorch RMSprop 的直接复制粘贴,在 sqrt 之前应用了 eps,并进行了其他一些修改,以更接近 Tensorflow,以便匹配超参数。

值得注意的更改包括:

  1. 平方根内部应用的 Epsilon
  2. square_avg 初始化为全 1
  3. 动量缓冲区中累积更新的 LR 缩放

由 G. Hinton 在其课程中提出。

居中版本首次出现在使用循环神经网络生成序列中。

步骤

< >

( closure = None )

参数

  • closure(可调用对象,可选)— 一个闭包,用于重新评估模型并返回损失。

执行单个优化步骤。

类 timm.optim.SGDP

< >

( paramslr = torch.optim.optimizer.requiredmomentum = 0dampening = 0weight_decay = 0nesterov = Falseeps = 1e-08delta = 0.1wd_ratio = 0.1caution = False )

类 timm.optim.SGDW

< >

( params: ParamsTlr: float = 0.001momentum: float = 0.0dampening: float = 0.0weight_decay: float = 0.0nesterov: bool = Falsecaution: bool = Falsecorrected_weight_decay: bool = Falsemaximize: bool = Falseforeach: typing.Optional[bool] = Nonedifferentiable: bool = False )

步骤

< >

( closure = None )

参数

  • closure (Callable, optional) — 一个闭包,用于重新评估模型并返回损失。

执行单个优化步骤。

Update on GitHub