Skip to content

schedulers

sleap_nn.training.schedulers

Custom learning rate schedulers for sleap-nn training.

This module provides learning rate schedulers with warmup phases that are commonly used in deep learning for pose estimation and computer vision tasks.

Classes:

Name Description
LinearWarmupCosineAnnealingLR

Cosine annealing learning rate scheduler with linear warmup.
LinearWarmupLinearDecayLR

Linear warmup followed by linear decay learning rate scheduler.

LinearWarmupCosineAnnealingLR

Bases: LRScheduler

Cosine annealing learning rate scheduler with linear warmup.

The learning rate increases linearly from warmup_start_lr to the optimizer's base learning rate over warmup_epochs, then decreases following a cosine curve to eta_min over the remaining epochs.

This schedule is widely used in vision transformers and modern CNN architectures as it provides stable early training (warmup) and smooth convergence (cosine decay).

Parameters:

Name Type Description Default
optimizer

Wrapped optimizer.

 required
warmup_epochs int

Number of epochs for the linear warmup phase.

 required
max_epochs int

Total number of training epochs.

 required
warmup_start_lr float

Learning rate at the start of warmup. Default: 0.0.

 0.0
eta_min float

Minimum learning rate at the end of the schedule. Default: 0.0.

 0.0
last_epoch int

The index of the last epoch. Default: -1.

 -1
Example

optimizer = torch.optim.Adam(model.parameters(), lr=1e-3) scheduler = LinearWarmupCosineAnnealingLR( ... optimizer, warmup_epochs=5, max_epochs=100, eta_min=1e-6 ... ) for epoch in range(100): ... train(...) ... scheduler.step()

Methods:

Name Description
__init__

Initialize the scheduler.
get_lr

Compute the learning rate at the current epoch.
Source code in sleap_nn/training/schedulers.py 
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
class LinearWarmupCosineAnnealingLR(LRScheduler):
    """Cosine annealing learning rate scheduler with linear warmup.

    The learning rate increases linearly from `warmup_start_lr` to the optimizer's
    base learning rate over `warmup_epochs`, then decreases following a cosine
    curve to `eta_min` over the remaining epochs.

    This schedule is widely used in vision transformers and modern CNN architectures
    as it provides stable early training (warmup) and smooth convergence (cosine decay).

    Args:
        optimizer: Wrapped optimizer.
        warmup_epochs: Number of epochs for the linear warmup phase.
        max_epochs: Total number of training epochs.
        warmup_start_lr: Learning rate at the start of warmup. Default: 0.0.
        eta_min: Minimum learning rate at the end of the schedule. Default: 0.0.
        last_epoch: The index of the last epoch. Default: -1.

    Example:
        >>> optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
        >>> scheduler = LinearWarmupCosineAnnealingLR(
        ...     optimizer, warmup_epochs=5, max_epochs=100, eta_min=1e-6
        ... )
        >>> for epoch in range(100):
        ...     train(...)
        ...     scheduler.step()
    """

    def __init__(
        self,
        optimizer,
        warmup_epochs: int,
        max_epochs: int,
        warmup_start_lr: float = 0.0,
        eta_min: float = 0.0,
        last_epoch: int = -1,
    ):
        """Initialize the scheduler.

        Args:
            optimizer: Wrapped optimizer.
            warmup_epochs: Number of epochs for the linear warmup phase.
            max_epochs: Total number of training epochs.
            warmup_start_lr: Learning rate at the start of warmup. Default: 0.0.
            eta_min: Minimum learning rate at the end of the schedule. Default: 0.0.
            last_epoch: The index of the last epoch. Default: -1.
        """
        if warmup_epochs < 0:
            raise ValueError(f"warmup_epochs must be >= 0, got {warmup_epochs}")
        if max_epochs <= 0:
            raise ValueError(f"max_epochs must be > 0, got {max_epochs}")
        if warmup_epochs >= max_epochs:
            raise ValueError(
                f"warmup_epochs ({warmup_epochs}) must be < max_epochs ({max_epochs})"
            )
        if warmup_start_lr < 0:
            raise ValueError(f"warmup_start_lr must be >= 0, got {warmup_start_lr}")
        if eta_min < 0:
            raise ValueError(f"eta_min must be >= 0, got {eta_min}")

        self.warmup_epochs = warmup_epochs
        self.max_epochs = max_epochs
        self.warmup_start_lr = warmup_start_lr
        self.eta_min = eta_min
        super().__init__(optimizer, last_epoch)

    def get_lr(self):
        """Compute the learning rate at the current epoch."""
        if self.last_epoch < self.warmup_epochs:
            # Linear warmup phase
            if self.warmup_epochs == 0:
                return list(self.base_lrs)
            alpha = self.last_epoch / self.warmup_epochs
            return [
                self.warmup_start_lr + alpha * (base_lr - self.warmup_start_lr)
                for base_lr in self.base_lrs
            ]
        else:
            # Cosine annealing phase
            decay_epochs = self.max_epochs - self.warmup_epochs
            if decay_epochs == 0:
                return [self.eta_min for _ in self.base_lrs]
            progress = (self.last_epoch - self.warmup_epochs) / decay_epochs
            # Clamp progress to [0, 1] to handle epochs beyond max_epochs
            progress = min(1.0, progress)
            return [
                self.eta_min
                + (base_lr - self.eta_min) * (1 + math.cos(math.pi * progress)) / 2
                for base_lr in self.base_lrs
            ]

__init__(optimizer, warmup_epochs, max_epochs, warmup_start_lr=0.0, eta_min=0.0, last_epoch=-1)

Initialize the scheduler.

Parameters:

Name Type Description Default
optimizer

Wrapped optimizer.

 required
warmup_epochs int

Number of epochs for the linear warmup phase.

 required
max_epochs int

Total number of training epochs.

 required
warmup_start_lr float

Learning rate at the start of warmup. Default: 0.0.

 0.0
eta_min float

Minimum learning rate at the end of the schedule. Default: 0.0.

 0.0
last_epoch int

The index of the last epoch. Default: -1.

 -1
Source code in sleap_nn/training/schedulers.py 
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
def __init__(
    self,
    optimizer,
    warmup_epochs: int,
    max_epochs: int,
    warmup_start_lr: float = 0.0,
    eta_min: float = 0.0,
    last_epoch: int = -1,
):
    """Initialize the scheduler.

    Args:
        optimizer: Wrapped optimizer.
        warmup_epochs: Number of epochs for the linear warmup phase.
        max_epochs: Total number of training epochs.
        warmup_start_lr: Learning rate at the start of warmup. Default: 0.0.
        eta_min: Minimum learning rate at the end of the schedule. Default: 0.0.
        last_epoch: The index of the last epoch. Default: -1.
    """
    if warmup_epochs < 0:
        raise ValueError(f"warmup_epochs must be >= 0, got {warmup_epochs}")
    if max_epochs <= 0:
        raise ValueError(f"max_epochs must be > 0, got {max_epochs}")
    if warmup_epochs >= max_epochs:
        raise ValueError(
            f"warmup_epochs ({warmup_epochs}) must be < max_epochs ({max_epochs})"
        )
    if warmup_start_lr < 0:
        raise ValueError(f"warmup_start_lr must be >= 0, got {warmup_start_lr}")
    if eta_min < 0:
        raise ValueError(f"eta_min must be >= 0, got {eta_min}")

    self.warmup_epochs = warmup_epochs
    self.max_epochs = max_epochs
    self.warmup_start_lr = warmup_start_lr
    self.eta_min = eta_min
    super().__init__(optimizer, last_epoch)

get_lr()

Compute the learning rate at the current epoch.

Source code in sleap_nn/training/schedulers.py 
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
def get_lr(self):
    """Compute the learning rate at the current epoch."""
    if self.last_epoch < self.warmup_epochs:
        # Linear warmup phase
        if self.warmup_epochs == 0:
            return list(self.base_lrs)
        alpha = self.last_epoch / self.warmup_epochs
        return [
            self.warmup_start_lr + alpha * (base_lr - self.warmup_start_lr)
            for base_lr in self.base_lrs
        ]
    else:
        # Cosine annealing phase
        decay_epochs = self.max_epochs - self.warmup_epochs
        if decay_epochs == 0:
            return [self.eta_min for _ in self.base_lrs]
        progress = (self.last_epoch - self.warmup_epochs) / decay_epochs
        # Clamp progress to [0, 1] to handle epochs beyond max_epochs
        progress = min(1.0, progress)
        return [
            self.eta_min
            + (base_lr - self.eta_min) * (1 + math.cos(math.pi * progress)) / 2
            for base_lr in self.base_lrs
        ]

LinearWarmupLinearDecayLR

Bases: LRScheduler

Linear warmup followed by linear decay learning rate scheduler.

The learning rate increases linearly from warmup_start_lr to the optimizer's base learning rate over warmup_epochs, then decreases linearly to end_lr over the remaining epochs.

This schedule provides a simple, interpretable learning rate trajectory and is commonly used in transformer-based models and NLP tasks.

Parameters:

Name Type Description Default
optimizer

Wrapped optimizer.

 required
warmup_epochs int

Number of epochs for the linear warmup phase.

 required
max_epochs int

Total number of training epochs.

 required
warmup_start_lr float

Learning rate at the start of warmup. Default: 0.0.

 0.0
end_lr float

Learning rate at the end of training. Default: 0.0.

 0.0
last_epoch int

The index of the last epoch. Default: -1.

 -1
Example

optimizer = torch.optim.Adam(model.parameters(), lr=1e-3) scheduler = LinearWarmupLinearDecayLR( ... optimizer, warmup_epochs=5, max_epochs=100, end_lr=1e-6 ... ) for epoch in range(100): ... train(...) ... scheduler.step()

Methods:

Name Description
__init__

Initialize the scheduler.
get_lr

Compute the learning rate at the current epoch.
Source code in sleap_nn/training/schedulers.py 
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
class LinearWarmupLinearDecayLR(LRScheduler):
    """Linear warmup followed by linear decay learning rate scheduler.

    The learning rate increases linearly from `warmup_start_lr` to the optimizer's
    base learning rate over `warmup_epochs`, then decreases linearly to `end_lr`
    over the remaining epochs.

    This schedule provides a simple, interpretable learning rate trajectory and is
    commonly used in transformer-based models and NLP tasks.

    Args:
        optimizer: Wrapped optimizer.
        warmup_epochs: Number of epochs for the linear warmup phase.
        max_epochs: Total number of training epochs.
        warmup_start_lr: Learning rate at the start of warmup. Default: 0.0.
        end_lr: Learning rate at the end of training. Default: 0.0.
        last_epoch: The index of the last epoch. Default: -1.

    Example:
        >>> optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
        >>> scheduler = LinearWarmupLinearDecayLR(
        ...     optimizer, warmup_epochs=5, max_epochs=100, end_lr=1e-6
        ... )
        >>> for epoch in range(100):
        ...     train(...)
        ...     scheduler.step()
    """

    def __init__(
        self,
        optimizer,
        warmup_epochs: int,
        max_epochs: int,
        warmup_start_lr: float = 0.0,
        end_lr: float = 0.0,
        last_epoch: int = -1,
    ):
        """Initialize the scheduler.

        Args:
            optimizer: Wrapped optimizer.
            warmup_epochs: Number of epochs for the linear warmup phase.
            max_epochs: Total number of training epochs.
            warmup_start_lr: Learning rate at the start of warmup. Default: 0.0.
            end_lr: Learning rate at the end of training. Default: 0.0.
            last_epoch: The index of the last epoch. Default: -1.
        """
        if warmup_epochs < 0:
            raise ValueError(f"warmup_epochs must be >= 0, got {warmup_epochs}")
        if max_epochs <= 0:
            raise ValueError(f"max_epochs must be > 0, got {max_epochs}")
        if warmup_epochs >= max_epochs:
            raise ValueError(
                f"warmup_epochs ({warmup_epochs}) must be < max_epochs ({max_epochs})"
            )
        if warmup_start_lr < 0:
            raise ValueError(f"warmup_start_lr must be >= 0, got {warmup_start_lr}")
        if end_lr < 0:
            raise ValueError(f"end_lr must be >= 0, got {end_lr}")

        self.warmup_epochs = warmup_epochs
        self.max_epochs = max_epochs
        self.warmup_start_lr = warmup_start_lr
        self.end_lr = end_lr
        super().__init__(optimizer, last_epoch)

    def get_lr(self):
        """Compute the learning rate at the current epoch."""
        if self.last_epoch < self.warmup_epochs:
            # Linear warmup phase
            if self.warmup_epochs == 0:
                return list(self.base_lrs)
            alpha = self.last_epoch / self.warmup_epochs
            return [
                self.warmup_start_lr + alpha * (base_lr - self.warmup_start_lr)
                for base_lr in self.base_lrs
            ]
        else:
            # Linear decay phase
            decay_epochs = self.max_epochs - self.warmup_epochs
            if decay_epochs == 0:
                return [self.end_lr for _ in self.base_lrs]
            progress = (self.last_epoch - self.warmup_epochs) / decay_epochs
            # Clamp progress to [0, 1] to handle epochs beyond max_epochs
            progress = min(1.0, progress)
            return [
                base_lr + progress * (self.end_lr - base_lr)
                for base_lr in self.base_lrs
            ]

__init__(optimizer, warmup_epochs, max_epochs, warmup_start_lr=0.0, end_lr=0.0, last_epoch=-1)

Initialize the scheduler.

Parameters:

Name Type Description Default
optimizer

Wrapped optimizer.

 required
warmup_epochs int

Number of epochs for the linear warmup phase.

 required
max_epochs int

Total number of training epochs.

 required
warmup_start_lr float

Learning rate at the start of warmup. Default: 0.0.

 0.0
end_lr float

Learning rate at the end of training. Default: 0.0.

 0.0
last_epoch int

The index of the last epoch. Default: -1.

 -1
Source code in sleap_nn/training/schedulers.py 
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
def __init__(
    self,
    optimizer,
    warmup_epochs: int,
    max_epochs: int,
    warmup_start_lr: float = 0.0,
    end_lr: float = 0.0,
    last_epoch: int = -1,
):
    """Initialize the scheduler.

    Args:
        optimizer: Wrapped optimizer.
        warmup_epochs: Number of epochs for the linear warmup phase.
        max_epochs: Total number of training epochs.
        warmup_start_lr: Learning rate at the start of warmup. Default: 0.0.
        end_lr: Learning rate at the end of training. Default: 0.0.
        last_epoch: The index of the last epoch. Default: -1.
    """
    if warmup_epochs < 0:
        raise ValueError(f"warmup_epochs must be >= 0, got {warmup_epochs}")
    if max_epochs <= 0:
        raise ValueError(f"max_epochs must be > 0, got {max_epochs}")
    if warmup_epochs >= max_epochs:
        raise ValueError(
            f"warmup_epochs ({warmup_epochs}) must be < max_epochs ({max_epochs})"
        )
    if warmup_start_lr < 0:
        raise ValueError(f"warmup_start_lr must be >= 0, got {warmup_start_lr}")
    if end_lr < 0:
        raise ValueError(f"end_lr must be >= 0, got {end_lr}")

    self.warmup_epochs = warmup_epochs
    self.max_epochs = max_epochs
    self.warmup_start_lr = warmup_start_lr
    self.end_lr = end_lr
    super().__init__(optimizer, last_epoch)

get_lr()

Compute the learning rate at the current epoch.

Source code in sleap_nn/training/schedulers.py 
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
def get_lr(self):
    """Compute the learning rate at the current epoch."""
    if self.last_epoch < self.warmup_epochs:
        # Linear warmup phase
        if self.warmup_epochs == 0:
            return list(self.base_lrs)
        alpha = self.last_epoch / self.warmup_epochs
        return [
            self.warmup_start_lr + alpha * (base_lr - self.warmup_start_lr)
            for base_lr in self.base_lrs
        ]
    else:
        # Linear decay phase
        decay_epochs = self.max_epochs - self.warmup_epochs
        if decay_epochs == 0:
            return [self.end_lr for _ in self.base_lrs]
        progress = (self.last_epoch - self.warmup_epochs) / decay_epochs
        # Clamp progress to [0, 1] to handle epochs beyond max_epochs
        progress = min(1.0, progress)
        return [
            base_lr + progress * (self.end_lr - base_lr)
            for base_lr in self.base_lrs
        ]