Spaces:
Running on Zero
Running on Zero
| """AdamW and step-addressed warmup/drop schedule.""" | |
| import torch | |
| def build_generator_optimizer(config, generator): | |
| settings = config.optimizer_config.adamw | |
| return torch.optim.AdamW( | |
| [parameter for parameter in generator.parameters() if parameter.requires_grad], | |
| lr=float(settings.lr), | |
| betas=(float(settings.beta1), float(settings.beta2)), | |
| weight_decay=float(settings.weight_decay), | |
| ) | |
| class LinearLrWarmup: | |
| """Updates 1..N use 1/N..1; drop nodes address completed optimizer updates.""" | |
| def __init__(self, optimizer, warmup_steps, drop_steps=(), drop_gamma=0.1): | |
| self.optimizer = optimizer | |
| self.warmup_steps = int(warmup_steps) | |
| self.drop_steps = tuple(drop_steps) | |
| self.drop_gamma = float(drop_gamma) | |
| self.base_lrs = tuple(float(group["lr"]) for group in optimizer.param_groups) | |
| def apply(self, step): | |
| if self.warmup_steps == 0 and not self.drop_steps: | |
| return 1.0 | |
| update = step + 1 | |
| factor = min(update / self.warmup_steps, 1.0) if self.warmup_steps else 1.0 | |
| factor *= self.drop_gamma ** sum(update >= node for node in self.drop_steps) | |
| for group, base_lr in zip(self.optimizer.param_groups, self.base_lrs, strict=True): | |
| group["lr"] = base_lr * factor | |
| return factor | |
| def optimizer_learning_rate_metrics(optimizer): | |
| metrics = {"learning_rate": float(optimizer.param_groups[0]["lr"])} | |
| for index, group in enumerate(optimizer.param_groups[1:], start=1): | |
| metrics[f"learning_rate/group_{index}"] = float(group["lr"]) | |
| return metrics | |