Overview

Optimization helpers that torch.optim does not ship.

Modules:

bn_momentum –

Stepwise exponential decay of the BatchNorm momentum over the epochs.
param_groups –

Optimizer parameter groups: weight-decay rules (none on biases, other 1-D parameters or named parameters)
schedulers –

Learning-rate schedules with a linear warm-up, complementing torch.optim.lr_scheduler.

Classes:

BNMomentumScheduler –

Apply bn_momentum to a model at the start of each epoch.
CosineWarmupLR –

Cosine decay with linear warm-up (one cycle).
PolyLR –

Polynomial decay with linear warm-up (one cycle).

Functions:

set_bn_momentum –

Set the momentum of every BatchNorm1d / 2d / 3d layer of model.
generate_param_groups –

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

BNMomentumScheduler [source] ¶

BNMomentumScheduler(
    model: Module,
    *,
    init: float = 0.5,
    decay_rate: float = 0.5,
    decay_step: int = 20,
    clip: float = 0.001,
)

Apply bn_momentum to a model at the start of each epoch.

Parameters:

model (Module) –

Model whose BatchNorm layers are scheduled.
init (float, default: 0.5 ) –

Initial momentum \(m_0\).
decay_rate (float, default: 0.5 ) –

Per-step multiplicative decay \(\gamma\).
decay_step (int, default: 20 ) –

Epochs between decay steps \(s\).
clip (float, default: 0.001 ) –

Lower bound \(m_\text{clip}\).

Example

import torch.nn as nn

from torch_pointcloud.optim import BNMomentumScheduler

model = nn.Sequential(nn.Linear(4, 4), nn.BatchNorm1d(4))
scheduler = BNMomentumScheduler(model)
scheduler.step(epoch=20)  # 0.25
model[1].momentum  # 0.25

Methods:

step –

Set the momentum of epoch on the model's BatchNorm layers and return it.

step [source] ¶

step(epoch: int) -> float

Set the momentum of epoch on the model's BatchNorm layers and return it.

CosineWarmupLR [source] ¶

CosineWarmupLR(
    optimizer: Optimizer,
    total_steps: int,
    *,
    warmup_steps: int = 0,
    warmup_lr_init: float = 0.0,
    lr_min: float = 0.0,
    warmup_prefix: bool = False,
    last_epoch: int = -1,
)

Bases: _WarmupSchedule

Cosine decay with linear warm-up (one cycle).

\[ \eta_t = \eta_\min + \tfrac{1}{2} (\eta_\text{base} - \eta_\min) \left(1 + \cos \frac{\pi\, t}{T}\right) \]

after a linear warm-up from warmup_lr_init over warmup_steps. With warmup_prefix=False (the default), \(t\) counts from step \(0\) and \(T\) is total_steps, so the warm-up takes its steps from the cosine; with warmup_prefix=True, \(t\) and \(T\) count from the end of the warm-up.

Parameters:

optimizer (Optimizer) –

Wrapped optimizer.
total_steps (int) –

Length of the schedule in steps (epochs or iterations, whichever step() counts).
warmup_steps (int, default: 0 ) –

Length of the linear warm-up.
warmup_lr_init (float, default: 0.0 ) –

Learning rate at step \(0\) of the warm-up.
lr_min (float, default: 0.0 ) –

Final (and post-schedule) learning rate.
warmup_prefix (bool, default: False ) –

Count the cosine from the end of the warm-up instead of from step \(0\).
last_epoch (int, default: -1 ) –

The index of the last step, for resuming.

Example

import torch

from torch_pointcloud.optim import CosineWarmupLR

opt = torch.optim.SGD([torch.zeros(1, requires_grad=True)], lr=1.0)
sched = CosineWarmupLR(opt, total_steps=10, warmup_steps=2, lr_min=0.0)
lrs = []
for _ in range(10):
    opt.step()
    sched.step()
    lrs.append(round(sched.get_last_lr()[0], 3))
lrs  # [0.5, 0.905, 0.794, 0.655, 0.5, 0.345, 0.206, 0.095, 0.024, 0.0]

PolyLR [source] ¶

PolyLR(
    optimizer: Optimizer,
    total_steps: int,
    *,
    power: float = 0.9,
    warmup_steps: int = 0,
    warmup_lr_init: float = 0.0,
    lr_min: float = 0.0,
    warmup_prefix: bool = False,
    last_epoch: int = -1,
)

Bases: _WarmupSchedule

Polynomial decay with linear warm-up (one cycle).

\[ \eta_t = \eta_\min + (\eta_\text{base} - \eta_\min) \left(1 - \frac{t}{T}\right)^{p} \]

after the warm-up; see CosineWarmupLR for how \(t\) and \(T\) are counted. The decay reaches lr_min at step total_steps; to keep a non-zero rate on the last of \(n\) steps, pass total_steps=n + 1.

Parameters:

optimizer (Optimizer) –

Wrapped optimizer.
total_steps (int) –

Length of the schedule in steps.
power (float, default: 0.9 ) –

Exponent \(p\) of the decay.
warmup_steps (int, default: 0 ) –

Length of the linear warm-up.
warmup_lr_init (float, default: 0.0 ) –

Learning rate at step \(0\) of the warm-up.
lr_min (float, default: 0.0 ) –

Final (and post-schedule) learning rate.
warmup_prefix (bool, default: False ) –

Count the decay from the end of the warm-up instead of from step \(0\).
last_epoch (int, default: -1 ) –

The index of the last step, for resuming.

Example

import torch

from torch_pointcloud.optim import PolyLR

opt = torch.optim.SGD([torch.zeros(1, requires_grad=True)], lr=1.0)
sched = PolyLR(opt, total_steps=4, power=1.0)
lrs = []
for _ in range(4):
    opt.step()
    sched.step()
    lrs.append(round(sched.get_last_lr()[0], 3))
lrs  # [0.75, 0.5, 0.25, 0.0]

set_bn_momentum [source] ¶

set_bn_momentum(model: Module, momentum: float) -> None

Set the momentum of every BatchNorm1d / 2d / 3d layer of model.

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