param_groups

Optimizer parameter groups: weight-decay rules (none on biases, other 1-D parameters or named parameters) and per-keyword overrides of the optimizer options, such as a lower learning rate for one sub-module.

Functions:

generate_param_groups –

Split a model's trainable parameters into optimizer groups with weight-decay rules and keyword overrides.

generate_param_groups [source] ¶

generate_param_groups(
    model: Module,
    *,
    weight_decay: float = 0.0,
    no_decay: NoDecayRule = False,
    overrides: Sequence[Mapping[str, Any]] = (),
) -> List[Dict[str, Any]]

Split a model's trainable parameters into optimizer groups with weight-decay rules and keyword overrides.

Every parameter with requires_grad goes to the first override whose keyword is a substring of its name (or whose match(name, param) returns True), else to the base group. Each group then carries weight_decay and the override's other options (lr, momentum, ...). With no_decay, every group is further split into a decaying and a non-decaying part: True excludes biases and 1-D parameters, a list of keywords extends that rule to the names containing them, and a callable decides per (name, param). An override that sets its own weight_decay is not split.

The groups come out base first, then the overrides in order, each as (decay, no-decay) when split; empty groups are dropped. That order is what a per-group max_lr list of OneCycleLR refers to.

Parameters:

model (Module) –

Source model.
weight_decay (float, default: 0.0 ) –

Weight decay of the base group (and of overrides that do not set their own).
no_decay (NoDecayRule, default: False ) –

Rule selecting the parameters that get no weight decay (see above).
overrides (Sequence[Mapping[str, Any]], default: () ) –

Dicts, each with keyword (substring of the parameter name) or match (callable (name, param) -> bool) plus the optimizer options of that group.

Returns:

List[Dict[str, Any]] –

Optimizer parameter-group dicts, each with params, weight_decay and the override options.

Example

>>> import torch.nn as nn
>>> model = nn.Sequential(nn.Linear(4, 4), nn.LayerNorm(4), nn.Linear(4, 2))
>>> groups = generate_param_groups(model, weight_decay=0.05, no_decay=True)
>>> [(len(g["params"]), g["weight_decay"]) for g in groups]
[(2, 0.05), (4, 0.0)]
>>> groups = generate_param_groups(model, weight_decay=0.05, overrides=[dict(keyword="2.", lr=1e-4)])
>>> [(len(g["params"]), g.get("lr")) for g in groups]
[(4, None), (2, 0.0001)]