File size: 801 Bytes
ad16788 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 |
"""PyTorch optimizer schdulers."""
from typing import List
from torch.optim import Optimizer
from espnet.scheduler.scheduler import SchedulerInterface
class PyTorchScheduler:
"""PyTorch optimizer scheduler."""
def __init__(self, schedulers: List[SchedulerInterface], optimizer: Optimizer):
"""Initialize class."""
self.schedulers = schedulers
self.optimizer = optimizer
for s in self.schedulers:
for group in optimizer.param_groups:
group.setdefault("initial_" + s.key, group[s.key])
def step(self, n_iter: int):
"""Update optimizer by scheduling."""
for s in self.schedulers:
for group in self.optimizer.param_groups:
group[s.key] = group["initial_" + s.key] * s.scale(n_iter)
|