Spaces:
Runtime error
Runtime error
| #!/usr/bin/env python3 | |
| # Copyright 2023 (authors: Feiteng Li) | |
| # | |
| # See ../../../../LICENSE for clarification regarding multiple authors | |
| # | |
| # Licensed under the Apache License, Version 2.0 (the "License"); | |
| # you may not use this file except in compliance with the License. | |
| # You may obtain a copy of the License at | |
| # | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, software | |
| # distributed under the License is distributed on an "AS IS" BASIS, | |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
| # See the License for the specific language governing permissions and | |
| # limitations under the License. | |
| import torch | |
| from modules.optim import Eden | |
| def calc_lr(step, dim_embed, warmup_steps): | |
| return dim_embed ** (-0.5) * min( | |
| step ** (-0.5), step * warmup_steps ** (-1.5) | |
| ) | |
| class NoamScheduler(torch.optim.lr_scheduler._LRScheduler): | |
| def __init__( | |
| self, | |
| base_lr: float, | |
| optimizer: torch.optim.Optimizer, | |
| dim_embed: int, | |
| warmup_steps: int, | |
| last_epoch: int = -1, | |
| verbose: bool = False, | |
| ) -> None: | |
| self.dim_embed = dim_embed | |
| self.base_lr = base_lr | |
| self.warmup_steps = warmup_steps | |
| self.num_param_groups = len(optimizer.param_groups) | |
| super().__init__(optimizer, last_epoch, verbose) | |
| def get_lr(self) -> float: | |
| lr = self.base_lr * calc_lr( | |
| self._step_count, self.dim_embed, self.warmup_steps | |
| ) | |
| return [lr] * self.num_param_groups | |
| def set_step(self, step: int): | |
| self._step_count = step | |
| def get_scheduler(params, optimizer): | |
| if params.scheduler_name.lower() == "eden": | |
| scheduler = Eden(optimizer, 5000, 4, warmup_batches=params.warmup_steps) | |
| elif params.scheduler_name.lower() == "noam": | |
| scheduler = NoamScheduler( | |
| params.base_lr, | |
| optimizer, | |
| params.decoder_dim, | |
| warmup_steps=params.warmup_steps, | |
| ) | |
| # scheduler.set_step(params.start_batch or params.batch_idx_train) | |
| elif params.scheduler_name.lower() == "cosine": | |
| scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( | |
| params.warmup_steps, | |
| optimizer, | |
| eta_min=params.base_lr, | |
| ) | |
| else: | |
| raise NotImplementedError(f"{params.scheduler_name}") | |
| return scheduler | |