"""PyTorch optimizer schdulers."""
from typing import List
from torch.optim import Optimizer
from espnet.scheduler.scheduler import SchedulerInterface
[docs]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])
[docs] 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)