bn_momentum

Stepwise exponential decay of the BatchNorm momentum over the epochs.

Classes:

BNMomentumScheduler –

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

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.

bn_momentum [source] ¶

bn_momentum(
    epoch: int,
    *,
    init: float = 0.5,
    decay_rate: float = 0.5,
    decay_step: int = 20,
    clip: float = 0.001,
) -> float

The decayed BatchNorm momentum of an epoch.

\[ m_\text{epoch} = \max\left(m_0 \cdot \gamma^{\lfloor \text{epoch} / s \rfloor},\; m_\text{clip}\right) \]

Parameters:

epoch (int) –

Zero-based epoch.
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}\).

Returns:

float –

The momentum to set on the BatchNorm layers for that epoch.

Example

>>> [bn_momentum(e) for e in (0, 19, 20, 40, 200)]
[0.5, 0.5, 0.25, 0.125, 0.001]

set_bn_momentum [source] ¶

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

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