Spaces:
Running on Zero
Running on Zero
File size: 1,619 Bytes
05182f3 | 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 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 | """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
|