pytorch-image-models

1
""" Step Scheduler
2

3
Basic step LR schedule with warmup, noise.
4

5
Hacked together by / Copyright 2020 Ross Wightman
6
"""
7
import math
8
import torch
9
from typing import List
10

11

12
from .scheduler import Scheduler
13

14

15
class StepLRScheduler(Scheduler):
16
    """
17
    """
18

19
    def __init__(
20
            self,
21
            optimizer: torch.optim.Optimizer,
22
            decay_t: float,
23
            decay_rate: float = 1.,
24
            warmup_t=0,
25
            warmup_lr_init=0,
26
            warmup_prefix=True,
27
            t_in_epochs=True,
28
            noise_range_t=None,
29
            noise_pct=0.67,
30
            noise_std=1.0,
31
            noise_seed=42,
32
            initialize=True,
33
    ) -> None:
34
        super().__init__(
35
            optimizer,
36
            param_group_field="lr",
37
            t_in_epochs=t_in_epochs,
38
            noise_range_t=noise_range_t,
39
            noise_pct=noise_pct,
40
            noise_std=noise_std,
41
            noise_seed=noise_seed,
42
            initialize=initialize,
43
        )
44

45
        self.decay_t = decay_t
46
        self.decay_rate = decay_rate
47
        self.warmup_t = warmup_t
48
        self.warmup_lr_init = warmup_lr_init
49
        self.warmup_prefix = warmup_prefix
50
        if self.warmup_t:
51
            self.warmup_steps = [(v - warmup_lr_init) / self.warmup_t for v in self.base_values]
52
            super().update_groups(self.warmup_lr_init)
53
        else:
54
            self.warmup_steps = [1 for _ in self.base_values]
55

56
    def _get_lr(self, t: int) -> List[float]:
57
        if t < self.warmup_t:
58
            lrs = [self.warmup_lr_init + t * s for s in self.warmup_steps]
59
        else:
60
            if self.warmup_prefix:
61
                t = t - self.warmup_t
62
            lrs = [v * (self.decay_rate ** (t // self.decay_t)) for v in self.base_values]
63
        return lrs
64