167 lines
6.7 KiB
Python
167 lines
6.7 KiB
Python
import math
|
|
import torch
|
|
import warnings
|
|
|
|
from torch import optim
|
|
from contextlib import contextmanager
|
|
|
|
def append_dims(x, target_dims):
|
|
"""Appends dimensions to the end of a tensor until it has target_dims dimensions."""
|
|
dims_to_append = target_dims - x.ndim
|
|
if dims_to_append < 0:
|
|
raise ValueError(f'input has {x.ndim} dims but target_dims is {target_dims}, which is less')
|
|
return x[(...,) + (None,) * dims_to_append]
|
|
|
|
|
|
def n_params(module):
|
|
"""Returns the number of trainable parameters in a module."""
|
|
return sum(p.numel() for p in module.parameters())
|
|
|
|
|
|
@contextmanager
|
|
def train_mode(model, mode=True):
|
|
"""A context manager that places a model into training mode and restores
|
|
the previous mode on exit."""
|
|
modes = [module.training for module in model.modules()]
|
|
try:
|
|
yield model.train(mode)
|
|
finally:
|
|
for i, module in enumerate(model.modules()):
|
|
module.training = modes[i]
|
|
|
|
|
|
def eval_mode(model):
|
|
"""A context manager that places a model into evaluation mode and restores
|
|
the previous mode on exit."""
|
|
return train_mode(model, False)
|
|
|
|
|
|
@torch.no_grad()
|
|
def ema_update(model, averaged_model, decay):
|
|
"""Incorporates updated model parameters into an exponential moving averaged
|
|
version of a model. It should be called after each optimizer step."""
|
|
model_params = dict(model.named_parameters())
|
|
averaged_params = dict(averaged_model.named_parameters())
|
|
assert model_params.keys() == averaged_params.keys()
|
|
|
|
for name, param in model_params.items():
|
|
averaged_params[name].mul_(decay).add_(param, alpha=1 - decay)
|
|
|
|
model_buffers = dict(model.named_buffers())
|
|
averaged_buffers = dict(averaged_model.named_buffers())
|
|
assert model_buffers.keys() == averaged_buffers.keys()
|
|
|
|
for name, buf in model_buffers.items():
|
|
averaged_buffers[name].copy_(buf)
|
|
|
|
|
|
class EMAWarmup:
|
|
"""Implements an EMA warmup using an inverse decay schedule.
|
|
If inv_gamma=1 and power=1, implements a simple average. inv_gamma=1, power=2/3 are
|
|
good values for models you plan to train for a million or more steps (reaches decay
|
|
factor 0.999 at 31.6K steps, 0.9999 at 1M steps), inv_gamma=1, power=3/4 for models
|
|
you plan to train for less (reaches decay factor 0.999 at 10K steps, 0.9999 at
|
|
215.4k steps).
|
|
Args:
|
|
inv_gamma (float): Inverse multiplicative factor of EMA warmup. Default: 1.
|
|
power (float): Exponential factor of EMA warmup. Default: 1.
|
|
min_value (float): The minimum EMA decay rate. Default: 0.
|
|
max_value (float): The maximum EMA decay rate. Default: 1.
|
|
start_at (int): The epoch to start averaging at. Default: 0.
|
|
last_epoch (int): The index of last epoch. Default: 0.
|
|
"""
|
|
|
|
def __init__(self, inv_gamma=1., power=1., min_value=0., max_value=1., start_at=0,
|
|
last_epoch=0):
|
|
self.inv_gamma = inv_gamma
|
|
self.power = power
|
|
self.min_value = min_value
|
|
self.max_value = max_value
|
|
self.start_at = start_at
|
|
self.last_epoch = last_epoch
|
|
|
|
def state_dict(self):
|
|
"""Returns the state of the class as a :class:`dict`."""
|
|
return dict(self.__dict__.items())
|
|
|
|
def load_state_dict(self, state_dict):
|
|
"""Loads the class's state.
|
|
Args:
|
|
state_dict (dict): scaler state. Should be an object returned
|
|
from a call to :meth:`state_dict`.
|
|
"""
|
|
self.__dict__.update(state_dict)
|
|
|
|
def get_value(self):
|
|
"""Gets the current EMA decay rate."""
|
|
epoch = max(0, self.last_epoch - self.start_at)
|
|
value = 1 - (1 + epoch / self.inv_gamma) ** -self.power
|
|
return 0. if epoch < 0 else min(self.max_value, max(self.min_value, value))
|
|
|
|
def step(self):
|
|
"""Updates the step count."""
|
|
self.last_epoch += 1
|
|
|
|
|
|
class InverseLR(optim.lr_scheduler._LRScheduler):
|
|
"""Implements an inverse decay learning rate schedule with an optional exponential
|
|
warmup. When last_epoch=-1, sets initial lr as lr.
|
|
inv_gamma is the number of steps/epochs required for the learning rate to decay to
|
|
(1 / 2)**power of its original value.
|
|
Args:
|
|
optimizer (Optimizer): Wrapped optimizer.
|
|
inv_gamma (float): Inverse multiplicative factor of learning rate decay. Default: 1.
|
|
power (float): Exponential factor of learning rate decay. Default: 1.
|
|
warmup (float): Exponential warmup factor (0 <= warmup < 1, 0 to disable)
|
|
Default: 0.
|
|
final_lr (float): The final learning rate. Default: 0.
|
|
last_epoch (int): The index of last epoch. Default: -1.
|
|
verbose (bool): If ``True``, prints a message to stdout for
|
|
each update. Default: ``False``.
|
|
"""
|
|
|
|
def __init__(self, optimizer, inv_gamma=1., power=1., warmup=0., final_lr=0.,
|
|
last_epoch=-1, verbose=False):
|
|
self.inv_gamma = inv_gamma
|
|
self.power = power
|
|
if not 0. <= warmup < 1:
|
|
raise ValueError('Invalid value for warmup')
|
|
self.warmup = warmup
|
|
self.final_lr = final_lr
|
|
super().__init__(optimizer, last_epoch, verbose)
|
|
|
|
def get_lr(self):
|
|
if not self._get_lr_called_within_step:
|
|
warnings.warn("To get the last learning rate computed by the scheduler, "
|
|
"please use `get_last_lr()`.")
|
|
|
|
return self._get_closed_form_lr()
|
|
|
|
def _get_closed_form_lr(self):
|
|
warmup = 1 - self.warmup ** (self.last_epoch + 1)
|
|
lr_mult = (1 + self.last_epoch / self.inv_gamma) ** -self.power
|
|
return [warmup * max(self.final_lr, base_lr * lr_mult)
|
|
for base_lr in self.base_lrs]
|
|
|
|
|
|
def rand_log_normal(shape, loc=0., scale=1., device='cpu', dtype=torch.float32):
|
|
"""Draws samples from an lognormal distribution."""
|
|
return (torch.randn(shape, device=device, dtype=dtype) * scale + loc).exp()
|
|
|
|
|
|
def rand_log_logistic(shape, loc=0., scale=1., min_value=0., max_value=float('inf'), device='cpu', dtype=torch.float32):
|
|
"""Draws samples from an optionally truncated log-logistic distribution."""
|
|
min_value = torch.as_tensor(min_value, device=device, dtype=torch.float64)
|
|
max_value = torch.as_tensor(max_value, device=device, dtype=torch.float64)
|
|
min_cdf = min_value.log().sub(loc).div(scale).sigmoid()
|
|
max_cdf = max_value.log().sub(loc).div(scale).sigmoid()
|
|
u = torch.rand(shape, device=device, dtype=torch.float64) * (max_cdf - min_cdf) + min_cdf
|
|
return u.logit().mul(scale).add(loc).exp().to(dtype)
|
|
|
|
|
|
def rand_log_uniform(shape, min_value, max_value, device='cpu', dtype=torch.float32):
|
|
"""Draws samples from an log-uniform distribution."""
|
|
min_value = math.log(min_value)
|
|
max_value = math.log(max_value)
|
|
return (torch.rand(shape, device=device, dtype=dtype) * (max_value - min_value) + min_value).exp()
|