Hugging Face's logo

timm documentation

Optimization

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v1.0.30).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

Optimization

This page contains the API reference documentation for the optimizers included in timm, and an overview of how they are selected and configured through the optimizer factory.

Usage

Optimizers are created by name with create_optimizer_v2. Given a model, the factory builds the parameter groups (weight decay exclusions, optional layer-wise learning rate decay) before constructing the optimizer. Optimizer specific arguments pass through as keyword arguments.

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))

The training script (train.py) exposes the same through --opt, --lr, --weight-decay, --momentum, --opt-eps, --opt-betas, --layer-decay and --opt-kwargs key=value ... for any other optimizer argument.

Naming conventions

Registered names are lower case, a base optimizer name plus modifiers:

ModifierMeaningExamples
c prefixCautious variant, the update is masked where its sign disagrees with the gradient (Cautious Optimizers)cadamw, clion, cmars, cadopt
w suffixDecoupled (AdamW style) weight decay, for optimizers that default to coupled L2 decaysgdw, lambw, adoptw, novogradw, rmsproptfw
c suffixCorrected weight decay, the decay is scaled by lr² / max_lr instead of lr (Why Gradients Rapidly Increase Near the End of Training)adamc, nadamc, sgdc, lionc, adoptc, kronc
n prefixNesterov momentum (SGD, LARS and Muon families)nesterov, nlars, nmuon, nadamuon
lookahead_ prefixWrap the optimizer in Lookaheadlookahead_adamw
fused prefixNVIDIA Apex fused implementation (requires apex)fusedadamw, fusedlamb
bnb prefixbitsandbytes implementation, 8bit suffix for 8-bit state (requires bitsandbytes)bnbadamw, bnbadamw8bit

A few names predate the conventions: lambc and larc are LAMB / LARS with trust ratio clipping (the corrected decay LAMB variant is lambcd), momentum is SGD with classical momentum and sgd is SGD with Nesterov momentum. adam, adamw, nadam, radam, radamw, adamax, adadelta, adagrad, rmsprop, sgd, momentum and nesterov pass through to torch.optim, the timm versions of AdamW / NAdam / RAdam that predate their inclusion in PyTorch are adamwlegacy, nadamlegacy and radamlegacy.

Optimizer families

NamesImplementationNotes
adamw, adam, nadam, radam, radamw, adamax, adadelta, adagrad, rmsprop, sgd, momentum, nesterovtorch.optim
adamwlegacy, adamc, cadamw, cadamcAdamWLegacyAdamW with caution, corrected decay, update clipping, capturable
nadamw, nadamc, cnadamw, cnadamcNAdamWAdamW with Nesterov momentum (MLCommons algorithmic-efficiency form)
adopt, adoptw, adoptc, cadopt, cadoptw, cadoptcAdoptAdam variant that converges for any β2
adan, adanw, cadan, cadanwAdanAdaptive Nesterov momentum
adabelief, radabeliefAdaBeliefradabelief adds RAdam style rectification
adafactor, cadafactorAdafactorfairseq Adafactor, relative step size when lr=None
adafactorbv, adafactorbvc, cadafactorbv, cadafactorbvcAdafactorBigVisionBig Vision Adafactor, factored second moments, low precision momentum, update clipping
adamp, cadamp, sgdp, csgdpAdamP, SGDPProjection of the update for scale invariant weights
adahessianAdahessianSecond order, requires a closure (is_second_order)
kron, kronw, kroncKronPSGD with a Kronecker-factored preconditioner
lamb, lambc, lambw, lambcw, lambcd, clamb, clambwLambLayer-wise adaptive moments, trust_clip variants (lambc, lambcw)
laprop, lapropc, clapropLaPropMomentum and adaptivity separated
lars, larc, nlars, nlarcLarsLayer-wise adaptive rate scaling, trust_clip variants (larc, nlarc)
lion, lionc, clion, clioncLionEvolved sign momentum
madgrad, madgradwMADGRADMomentum based dual averaging
mars, cmarsMarsVariance reduction, AdamW or Lion style update
muon, nmuon, adamuon, nadamuonMuonNewton-Schulz orthogonalized momentum for 2D+ weights, AdamW / NAdamW fallback for 1D params
novograd, novogradw, novogradcNvNovoGradLayer-wise normalized moments
rmsproptf, rmsproptfw, rmsproptfc, crmsproptfRMSpropTFTensorFlow ordering of ops and eps handling, lr folded into the momentum buffer
sgdw, nesterovw, sgdc, csgdw, csgdcSGDWSGD with decoupled weight decay
nadamlegacy, radamlegacyNAdamLegacy, RAdamLegacyReference implementations, superseded by torch.optim
fused*, bnb*NVIDIA Apex, bitsandbytesRegistered when the package is installed

Common options

Most timm optimizers share a set of options beyond the usual lr, weight_decay, betas and eps. Which optimizers accept which can be checked in the class signatures below, the factory raises on unsupported arguments.

  • caution applies the cautious update mask, equivalent to the c prefixed registry names.
  • corrected_weight_decay scales decoupled weight decay by lr² / max_lr instead of lr (the c suffixed names), so the effective decay follows the learning rate schedule. max_lr is the learning rate the optimizer was constructed with, construct it with the peak schedule lr (as train.py does). It is kept in each param group (max_lr_snapshot) and saved / restored with the optimizer state.
  • decoupled_decay (or decoupled) switches optimizers that default to coupled L2 decay to AdamW style decay (the w suffixed names).
  • foreach selects the multi-tensor implementation. By default (None) it is used on CUDA and other devices with foreach kernels, matching torch.optim, unless unsupported by the configuration (a tensor lr without capturable, or caution on a PyTorch version without the required _foreach_maximum_ overload). The multi-tensor step is run in chunks of parameters sized so each chunk’s working set stays in the GPU L2 cache, which is faster than streaming every parameter through memory once per op for large models. The chunk size is derived from the device L2 cache size (PyTorch >= 2.5) and chunking is skipped on devices where the resulting chunk would be below 4M elements. TIMM_FOREACH_CHUNK_SIZE=<elements> overrides the chunk size, 0 disables chunking.
  • capturable (AdamWLegacy, NAdamW, Adopt, Muon) keeps the step count on the device so the step can be captured in a CUDA graph. All optimizers accept an lr given as a 0-dim tensor, updated in place by a schedule, for use with torch.compile and CUDA graphs.
  • clipping_threshold (AdamWLegacy, NAdamW, AdafactorBigVision) clips the update by its RMS, update / max(1, rms / threshold) with rms the RMS of the second-moment normalized gradient, per tensor (StableAdamW, Adafactor). 1.0 as in the papers. As in Adafactor, the (decoupled) weight decay is not clipped, StableAdamW also scales the decay by the clipped learning rate.
  • maximize maximizes the objective instead of minimizing.

Low precision parameters (bfloat16, float16) are supported throughout. AdaBelief, RAdamLegacy and Adafactor keep their state in float32 for low precision params, AdafactorBigVision keeps its momentum in a reduced precision (momentum_dtype) and Kron has separate mu_dtype / precond_dtype options. These state dtypes are preserved by load_state_dict (PyTorch’s default loader casts state to the param dtype). The step functions avoid host syncs (data dependent Python branches on device tensors), AdamP / SGDP’s projection check is the exception.

Parameter groups

When given a model, create_optimizer_v2 groups parameters with param_groups_weight_decay (no decay for 1D params and the model’s no_weight_decay() set) or, with layer_decay set, param_groups_layer_decay, which assigns a learning rate scale of layer_decay ** (num_layers - layer_id) per layer using the model’s group_matcher(). Models without a group_matcher() fall back to auto_group_layers, which needs the model’s pretrained_cfg['classifier'] to separate the head. For hybrid optimizers (Muon) the fallback_list and fallback_no_weight_decay arguments route parameters to the fallback (AdamW) optimizer.

API

Factory functions

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 )

Parameters

  • model_or_params — A PyTorch model or an iterable of parameters/parameter groups. If a model is provided, parameters will be automatically extracted and grouped based on the other arguments.
  • opt — Name of the optimizer to create (e.g., ‘adam’, ‘adamw’, ‘sgd’). Use list_optimizers() to see available options.
  • lr — Learning rate. If None, will use the optimizer’s default.
  • weight_decay — Weight decay factor. Will be used to create param groups if model_or_params is a model.
  • momentum — Momentum factor for optimizers that support it. Only used if the chosen optimizer accepts a momentum parameter.
  • foreach — Enable/disable foreach (multi-tensor) implementation if available. If None, will use optimizer-specific defaults.
  • filter_bias_and_bn — If True, bias, norm layer parameters (all 1d params) will not have weight decay applied. Only used when model_or_params is a model and weight_decay > 0.
  • fallback_list — Collection of parameter name patterns to use fallback optimizer for hybrid optimizers (e.g., AdamW for Muon). Supports wildcard matching.
  • fallback_no_weight_decay — If True, params in model’s no_weight_decay() list will use fallback optimizer for hybrid optimizers (e.g., AdamW for Muon).
  • layer_decay — Optional layer-wise learning rate decay factor. If provided, learning rates will be scaled by layer_decay^(max_depth - layer_depth). Only used when model_or_params is a model.
  • param_group_fn — Optional function to create custom parameter groups. If provided, other parameter grouping options will be ignored.
  • **kwargs — Additional optimizer-specific arguments (e.g., betas for Adam).

Create an optimizer instance via timm registry.

Creates and configures an optimizer with appropriate parameter groups and settings. Supports automatic parameter group creation for weight decay and layer-wise learning rates, as well as custom parameter grouping.

Examples:

Basic usage with a model

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

SGD with momentum and weight decay

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

Adam with layer-wise learning rate decay

optimizer = create_optimizer_v2( … model, ‘adam’, lr=1e-3, layer_decay=0.7 … )

Custom parameter groups

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 … )

Note: Parameter group handling precedence:

  1. If param_group_fn is provided, it will be used exclusively
  2. If layer_decay is provided, layer-wise groups will be created
  3. If weight_decay > 0 and filter_bias_and_bn is True, weight decay groups will be created
  4. Otherwise, all parameters will be in a single group

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

Parameters

  • filter — Wildcard style filter string or list of filter strings (e.g., ‘adam’ for all Adam variants, or [‘adam’, ‘*8bit’] for Adam variants and 8-bit optimizers). Empty string means no filtering.
  • exclude_filters — Optional list of wildcard patterns to exclude. For example, [’8bit’, ‘fused’] would exclude 8-bit and fused implementations.
  • with_description — If True, returns tuples of (name, description) instead of just names. Descriptions provide brief explanations of optimizer characteristics.

Returns

If with_description is False

List of optimizer names as strings (e.g., [‘adam’, ‘adamw’, …]) If with_description is True: List of tuples of (name, description) (e.g., [(‘adam’, ‘Adaptive Moment…’), …])

List available optimizer names, optionally filtered.

List all registered optimizers, with optional filtering using wildcard patterns. Optimizers can be filtered using include and exclude patterns, and can optionally return descriptions with each optimizer name.

Examples:

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

Parameters

  • name — Name of the optimizer to retrieve (e.g., ‘adam’, ‘sgd’)
  • bind_defaults — If True, returns a partial function with default arguments from OptimInfo bound. If False, returns the raw optimizer class.

Returns

If bind_defaults is False

The optimizer class (e.g., torch.optim.Adam) If bind_defaults is True: A partial function with default arguments bound

Raises

ValueError

  • ValueError — If optimizer name is not found in registry

Get optimizer class by name with option to bind default arguments.

Retrieves the optimizer class or a partial function with default arguments bound. This allows direct instantiation of optimizers with their default configurations without going through the full factory.

Examples:

Get SGD with nesterov momentum default

SGD = get_optimizer_class(‘sgd’) # nesterov=True bound opt = SGD(model.parameters(), lr=0.1, momentum=0.9)

Get raw optimizer class

SGD = get_optimizer_class(‘sgd’) opt = SGD(model.parameters(), lr=1e-3, momentum=0.9)

timm.optim.get_optimizer_info

< >

( name: str )

Parameters

  • name — Name of the optimizer

Raises

ValueError

  • ValueError — If optimizer is not found

Get the OptimInfo for an optimizer.

Parameter group functions

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 )

Parameter groups for layer-wise lr decay & weight decay Based on BEiT: https://github.com/microsoft/unilm/blob/master/beit/optim_factory.py#L58

timm.optim.auto_group_layers

< >

( modellayers_per_group = 12num_groups = None )

Optimizer Classes

class 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 )

Parameters

  • params (iterable) — iterable of parameters to optimize or dicts defining parameter groups
  • lr (float, optional) — learning rate (default: 1e-3)
  • betas (Tuple[float, float], optional) — coefficients used for computing running averages of gradient and its square (default: (0.9, 0.999))
  • eps (float, optional) — term added to the denominator to improve numerical stability (default: 1e-16)
  • weight_decay (float, optional) — weight decay (L2 penalty) (default: 0)
  • amsgrad (boolean, optional) — whether to use the AMSGrad variant of this algorithm from the paper On the Convergence of Adam and Beyond_ (default: False)
  • decoupled_decay (boolean, optional) — (default: True) If set as True, then the optimizer uses decoupled weight decay as in AdamW
  • fixed_decay (boolean, optional) — (default: False) This is used when weightdecouple is set as True. When fixed_decay == True, the weight decay is performed as $W{new} = W{old} - W{old} \times decay$. When fixeddecay == False, the weight decay is performed as $W{new} = W{old} - W{old} \times decay \times lr$. Note that in this case, the weight decay ratio decreases with learning rate (lr).
  • rectify (boolean, optional) — (default: True) If set as True, then perform the rectified update similar to RAdam
  • degenerated_to_sgd (boolean, optional) (default —True) If set as True, then perform SGD update when variance of gradient is high
  • foreach (boolean, optional) — use the foreach (multi-tensor) impl, by default (None) used unless unsupported

Implements AdaBelief algorithm. Modified from Adam in PyTorch

reference: AdaBelief Optimizer, adapting stepsizes by the belief in observed gradients, NeurIPS 2020

For a complete table of recommended hyperparameters, see https://github.com/juntang-zhuang/Adabelief-Optimizer’ For example train/args for EfficientNet see these gists

step

< >

( closure = None )

Parameters

  • closure (callable, optional) — A closure that reevaluates the model and returns the loss.

Performs a single optimization step.

class 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 )

Implements Adafactor algorithm.

This implementation is based on: Adafactor: Adaptive Learning Rates with Sublinear Memory Cost (see https://arxiv.org/abs/1804.04235)

Note that this optimizer internally adjusts the learning rate depending on the scale_parameter, relative_step and warmup_init options.

A manual (external) learning rate schedule is used when lr is set (relative_step is then False). Note that in this implementation scale_parameter and warmup_init only apply to the relative step learning rate (lr=None), unlike fairseq / HF Transformers, the parameter RMS scaling is not applied to an external learning rate.

Ags: params: iterable of parameters to optimize or dicts defining parameter groups lr: external learning rate eps: regularization constants for square gradient and parameter scale respectively eps_scale: regularization constants for parameter scale respectively clip_threshold: threshold of root-mean-square of final gradient update decay_rate: coefficient used to compute running averages of square gradient beta1: coefficient used for computing running averages of gradient weight_decay: weight decay scale_parameter: if True, relative step learning rate is scaled by root-mean-square of parameter warmup_init: time-dependent learning rate computation depends on whether warm-up initialization is being used

step

< >

( closure = None )

Parameters

  • closure (callable, optional) — A closure that reevaluates the model and returns the loss.

Performs a single optimization step.

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 implementation of BigVision’s Adafactor variant with both single and multi tensor implementations.

Adapted from https://github.com/google-research/big_vision by Ross Wightman

class 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 )

Parameters

  • params (iterable) — iterable of parameters to optimize or dicts defining parameter groups
  • lr (float, optional) — learning rate (default: 0.1)
  • betas ((float, float), optional) — coefficients used for computing running averages of gradient and the squared hessian trace (default: (0.9, 0.999))
  • eps (float, optional) — term added to the denominator to improve numerical stability (default: 1e-8)
  • weight_decay (float, optional) — weight decay (L2 penalty) (default: 0.0)
  • hessian_power (float, optional) — exponent of the hessian trace (default: 1.0)
  • update_each (int, optional) — compute the hessian trace approximation only after this number of steps (to save time) (default: 1)
  • n_samples (int, optional) — how many times to sample z for the approximation of the hessian trace (default: 1)

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

get_params

< >

( )

Gets all parameters in all param_groups with gradients

set_hessian

< >

( )

Computes the Hutchinson approximation of the hessian trace and accumulates it for each trainable parameter.

step

< >

( closure = None )

Parameters

  • closure (callable, optional) — a closure that reevaluates the model and returns the loss (default — None)

Performs a single optimization step.

zero_hessian

< >

( )

Zeros out the accumulated hessian traces.

class timm.optim.AdamP

< >

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

class 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 )

Parameters

  • params — iterable of parameters to optimize or dicts defining parameter groups
  • lr — learning rate
  • betas — coefficients used for computing running averages of gradient and its square
  • eps — term added to the denominator to improve numerical stability
  • weight_decay — weight decay coefficient
  • amsgrad — whether to use the AMSGrad variant of this algorithm from the paper On the Convergence of Adam and Beyond
  • caution — apply caution when using AdamW
  • corrected_weight_decay — apply corrected weight decay (lr**2 / max_lr)
  • maximize — maximize the params based on the objective, instead of minimizing
  • 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 — whether this instance is safe to capture in a CUDA graph. Passing True can impair ungraphed performance, so if you don’t intend to graph capture this instance, leave it False
  • clipping_threshold — RMS update clipping threshold (StableAdamW, https://arxiv.org/abs/2304.13013), the update is scaled by 1 / max(1, rms / clipping_threshold) w/ rms the RMS of grad / (sqrt(v_hat) + eps). Damps the update of params w/ a sudden gradient increase relative to their second moment estimate. 1.0 as per the paper, None (default) disables. As in Adafactor, the decoupled weight decay is not clipped (StableAdamW also scales the decay by the clipped lr).

Implements AdamW algorithm.

NOTE: This impl has been deprecated in favour of torch.optim.AdamW and remains as a reference

References:

step

< >

( closure = None )

Parameters

  • closure (callable, optional) — A closure that reevaluates the model and returns the loss.

Performs a single optimization step.

class 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 )

Parameters

  • params — Iterable of parameters to optimize or dicts defining parameter groups.
  • lr — Learning rate.
  • betas — Coefficients used for first- and second-order moments.
  • eps — Term added to the denominator to improve numerical stability.
  • weight_decay — Decoupled weight decay (L2 penalty)
  • no_prox — How to perform the weight decay
  • caution — Enable caution from ‘Cautious Optimizers’
  • foreach — If True would use torch._foreach implementation. Faster but uses slightly more memory.

Implements a pytorch variant of Adan.

Adan was proposed in Adan: Adaptive Nesterov Momentum Algorithm for Faster Optimizing Deep Models https://arxiv.org/abs/2208.06677

step

< >

( closure = None )

Performs a single optimization step.

class 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: Modified Adam Can Converge with Any β2 with the Optimal Rate: https://arxiv.org/abs/2411.02853

step

< >

( closure = None )

Parameters

  • closure (Callable, optional) — A closure that reevaluates the model and returns the loss.

Perform a single optimization step.

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 )

Parameters

  • params — Iterable of parameters to optimize or dicts defining parameter groups.
  • lr — Learning rate.
  • momentum — Momentum parameter.
  • weight_decay — Weight decay.
  • preconditioner_update_probability — Probability of updating the preconditioner. If None, defaults to a schedule that anneals from 1.0 to 0.03 by 4000 steps.
  • max_size_triangular — Max size for dim’s preconditioner to be triangular.
  • min_ndim_triangular — Minimum number of dimensions a layer needs to have triangular preconditioners.
  • memory_save_mode — ‘one_diag’, ‘smart_one_diag’, or ‘all_diag’, None is default to set all preconditioners to be triangular, ‘one_diag’ sets the largest or last dim to be diagonal per layer, and ‘all_diag’ sets all preconditioners to be diagonal.
  • momentum_into_precond_update — whether to send momentum into preconditioner update instead of raw gradients.
  • mu_dtype — Dtype of the momentum accumulator.
  • precond_dtype — Dtype of the preconditioner.
  • decoupled_decay — AdamW style decoupled weight decay
  • corrected_weight_decay — apply corrected weight decay when using decoupled_decay (lr**2 / max_lr)
  • flatten — Flatten dimensions instead of fully relying on expressions for higher rank params
  • flatten_start_dim — Start of flatten range, defaults to 2. Seems good tradeoff for ConvNets.
  • flatten_end_dim — End of flatten range, defaults to -1.
  • stochastic_weight_decay — Enable random modulation of weight decay
  • deterministic — Deterministic behaviour across save / load (resume). FIXME slow, needs work

Implements PSGD Kron from https://github.com/lixilinx/psgd_torch.

class 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 )

Parameters

  • params — Iterable of parameters to optimize or dicts defining parameter groups.
  • lr — Learning rate
  • betas — Coefficients used for computing running averages of gradient and its norm.
  • eps — Term added to the denominator to improve numerical stability.
  • weight_decay — Weight decay
  • grad_averaging — Whether apply (1-beta2) to grad when calculating running averages of gradient.
  • max_grad_norm — Value used to clip global grad norm.
  • trust_clip — Enable LAMBC trust ratio clipping.
  • always_adapt — Apply adaptive learning rate to 0.0 weight decay parameter.
  • caution — Apply caution.
  • decoupled — apply decoupled weight decay
  • corrected_weight_decay — apply corrected weight decay (lr**2 / max_lr) when using decoupled_decay
  • foreach — use the foreach (multi-tensor) impl, by default (None) used unless unsupported

Implements a pure pytorch variant of FuseLAMB (NvLamb variant) optimizer from apex.optimizers.FusedLAMB reference: https://github.com/NVIDIA/DeepLearningExamples/blob/master/PyTorch/LanguageModeling/Transformer-XL/pytorch/lamb.py

LAMB was proposed in:

step

< >

( closure = None )

Parameters

  • closure (callable, optional) — A closure that reevaluates the model and returns the loss.

Performs a single optimization step.

class 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 )

Parameters

  • foreach — use the foreach (multi-tensor) impl, by default (None) used unless unsupported (see laprop())

LaProp Optimizer

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

step

< >

( closure = None )

Parameters

  • closure (callable, optional) — A closure that reevaluates the model and returns the loss.

Performs a single optimization step.

class timm.optim.Lars

< >

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

Parameters

  • params (iterable) — iterable of parameters to optimize or dicts defining parameter groups.
  • lr (float, optional) — learning rate (default: 1.0).
  • momentum (float, optional) — momentum factor (default: 0)
  • weight_decay (float, optional) — weight decay (L2 penalty) (default: 0)
  • dampening (float, optional) — dampening for momentum (default: 0)
  • nesterov (bool, optional) — enables Nesterov momentum (default: False)
  • trust_coeff (float) — trust coefficient for computing adaptive lr / trust_ratio (default: 0.001)
  • eps (float) — eps for division denominator (default: 1e-8)
  • trust_clip (bool) — enable LARC trust ratio clipping (default: False)
  • always_adapt (bool) — always apply LARS LR adapt, otherwise only when group weight_decay != 0 (default: False)
  • foreach (bool, optional) — use the foreach (multi-tensor) impl, by default (None) used unless unsupported

LARS for PyTorch

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

step

< >

( closure = None )

Parameters

  • closure (callable, optional) — A closure that reevaluates the model and returns the loss.

Performs a single optimization step.

class 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 )

Implements Lion algorithm.

step

< >

( closure = None )

Parameters

  • closure — A closure that reevaluates the model and returns the loss.

Performs a single optimization step.

class timm.optim.Lookahead

< >

( base_optimizeralpha = 0.5k = 6 )

class timm.optim.MADGRAD

< >

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

Parameters

  • params (iterable) — Iterable of parameters to optimize or dicts defining parameter groups.
  • lr (float) — Learning rate (default: 1e-2).
  • momentum (float) — Momentum value in the range [0,1) (default: 0.9).
  • weight_decay (float) — Weight decay, i.e. a L2 penalty (default: 0).
  • eps (float) — Term added to the denominator outside of the root operation to improve numerical stability. (default: 1e-6).

MADGRAD_: A Momentumized, Adaptive, Dual Averaged Gradient Method for Stochastic Optimization.

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

MADGRAD is a general purpose optimizer that can be used in place of SGD or Adam may converge faster and generalize better. Currently GPU-only. Typically, the same learning rate schedule that is used for SGD or Adam may be used. The overall learning rate is not comparable to either method and should be determined by a hyper-parameter sweep.

MADGRAD requires less weight decay than other methods, often as little as zero. Momentum values used for SGD or Adam’s beta1 should work here also.

On sparse problems both weight_decay and momentum should be set to 0.

step

< >

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

Parameters

  • closure (callable, optional) — A closure that reevaluates the model and returns the loss.

Performs a single optimization step.

class 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 Optimizer

Paper: MARS: Unleashing the Power of Variance Reduction for Training Large Models https://arxiv.org/abs/2411.10438

step

< >

( closure = None )

Parameters

  • closure (callable, optional) — A closure that reevaluates the model and returns the loss.

Performs a single optimization step.

class 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 - MomentUm Orthogonalized by Newton-schulz

Combines Muon for 2D+ parameters (weight matrices) with AdamW for 1D parameters (biases, norms) and parameter groups with ‘use_fallback=True’ set (or ‘use_muon=False’ for compatibility).

Supports two algorithms:

  • “muon”: Standard Muon algorithm with momentum + orthogonalization
  • “adamuon”: AdaMuon algorithm that adds element-wise second moment estimation to orthogonalized directions for Adam-like adaptive scaling

step

< >

( closure = None )

Performs a single optimization step.

class timm.optim.NAdamLegacy

< >

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

Parameters

  • params (iterable) — iterable of parameters to optimize or dicts defining parameter groups
  • lr (float, optional) — learning rate (default: 2e-3)
  • betas (Tuple[float, float], optional) — coefficients used for computing running averages of gradient and its square
  • eps (float, optional) — term added to the denominator to improve numerical stability (default: 1e-8)
  • weight_decay (float, optional) — weight decay (L2 penalty) (default: 0)
  • schedule_decay (float, optional) — momentum schedule decay (default: 4e-3)

Implements Nadam algorithm (a variant of Adam based on Nesterov momentum).

NOTE: This impl has been deprecated in favour of torch.optim.NAdam and remains as a reference

It has been proposed in Incorporating Nesterov Momentum into Adam__.

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

Originally taken from: https://github.com/pytorch/pytorch/pull/1408 NOTE: Has potential issues but does work well on some problems.

step

< >

( closure = None )

Parameters

  • closure (callable, optional) — A closure that reevaluates the model and returns the loss.

Performs a single optimization step.

class 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 )

Parameters

  • params — iterable of parameters to optimize or dicts defining parameter groups
  • lr — learning rate
  • betas — coefficients used for computing running averages of gradient and its square
  • eps — term added to the denominator to improve numerical stability
  • weight_decay — weight decay coefficient
  • caution — enable caution
  • corrected_weight_decay — apply corrected weight decay (lr**2 / max_lr)
  • maximize — maximize the params based on the objective, instead of minimizing
  • foreach — whether foreach implementation of optimizer is used. If unspecified (None), the foreach implementation is used on CUDA (devices w/ foreach kernels) since it is faster in general.
  • capturable — whether this instance is safe to capture in a CUDA graph
  • clipping_threshold — RMS update clipping threshold (StableAdamW, https://arxiv.org/abs/2304.13013), the update is scaled by 1 / max(1, rms / clipping_threshold) w/ rms the RMS of grad / (sqrt(v_hat) + eps). Damps the update of params w/ a sudden gradient increase relative to their second moment estimate. 1.0 as per the paper, None (default) disables. As in Adafactor, the decoupled weight decay is not clipped (StableAdamW also scales the decay by the clipped lr).

Implements NAdamW algorithm.

See Table 1 in https://arxiv.org/abs/1910.05446 for the implementation of the NAdam algorithm (there is also a comment in the code which highlights the only difference of NAdamW and AdamW).

For further details regarding the algorithm we refer to

step

< >

( closure = None )

Parameters

  • closure (callable, optional) — A closure that reevaluates the model and returns the loss.

Performs a single optimization step.

class timm.optim.NvNovoGrad

< >

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

Parameters

  • params (iterable) — iterable of parameters to optimize or dicts defining parameter groups
  • lr (float, optional) — learning rate (default: 1e-3)
  • betas (Tuple[float, float], optional) — coefficients used for computing running averages of gradient and its square (default: (0.95, 0.98))
  • eps (float, optional) — term added to the denominator to improve numerical stability (default: 1e-8)
  • weight_decay (float, optional) — weight decay (L2 penalty) (default: 0)
  • grad_averaging — gradient averaging
  • amsgrad (boolean, optional) — whether to use the AMSGrad variant of this algorithm from the paper On the Convergence of Adam and Beyond_ (default: False)
  • decoupled_decay — apply weight decay directly to the weights (lr * weight_decay per step) instead of adding it to the normalized gradient before the momentum update (as per the paper)
  • corrected_weight_decay — apply corrected weight decay (lr**2 / max_lr) when decoupled_decay is True

Implements Novograd algorithm.

step

< >

( closure = None )

Parameters

  • closure (callable, optional) — A closure that reevaluates the model
  • and returns the loss. —

Performs a single optimization step.

class timm.optim.RAdamLegacy

< >

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

PyTorch RAdam optimizer

NOTE: This impl has been deprecated in favour of torch.optim.AdamW and remains as a reference

class 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 )

Parameters

  • params — iterable of parameters to optimize or dicts defining parameter groups
  • lr — learning rate
  • momentum — momentum factor
  • alpha — smoothing (decay) constant
  • eps — term added to the denominator to improve numerical stability
  • centered — if True, compute the centered RMSProp, the gradient is normalized by an estimation of its variance
  • weight_decay — weight decay (L2 penalty) (default: 0)
  • decoupled_decay — decoupled weight decay as per https://arxiv.org/abs/1711.05101
  • corrected_weight_decay — apply corrected weight decay (lr**2 / max_lr) when decoupled_decay is True
  • lr_in_momentum — learning rate scaling is included in the momentum buffer update as per defaults in Tensorflow
  • caution — apply caution
  • foreach — use the foreach (multi-tensor) impl, by default (None) used unless unsupported (see rmsprop_tf())

Implements RMSprop algorithm (TensorFlow style epsilon)

NOTE: This is a direct cut-and-paste of PyTorch RMSprop with eps applied before sqrt and a few other modifications to closer match Tensorflow for matching hyper-params.

Noteworthy changes include:

  1. Epsilon applied inside square-root
  2. square_avg initialized to ones
  3. LR scaling of update accumulated in momentum buffer

Proposed by G. Hinton in his course.

The centered version first appears in Generating Sequences With Recurrent Neural Networks.

step

< >

( closure = None )

Parameters

  • closure (callable, optional) — A closure that reevaluates the model and returns the loss.

Performs a single optimization step.

class timm.optim.SGDP

< >

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

class 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 )

step

< >

( closure = None )

Parameters

  • closure (Callable, optional) — A closure that reevaluates the model and returns the loss.

Performs a single optimization step.

Update on GitHub

←Losses Learning Rate Schedulers→