|
from contextlib import contextmanager |
|
import hashlib |
|
import math |
|
from pathlib import Path |
|
import shutil |
|
import urllib |
|
import warnings |
|
|
|
from PIL import Image |
|
import torch |
|
from torch import nn, optim |
|
from torch.utils import data |
|
|
|
|
|
def hf_datasets_augs_helper(examples, transform, image_key, mode='RGB'): |
|
"""Apply passed in transforms for HuggingFace Datasets.""" |
|
images = [transform(image.convert(mode)) for image in examples[image_key]] |
|
return {image_key: images} |
|
|
|
|
|
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') |
|
expanded = x[(...,) + (None,) * dims_to_append] |
|
|
|
|
|
return expanded.detach().clone() if expanded.device.type == 'mps' else expanded |
|
|
|
|
|
def n_params(module): |
|
"""Returns the number of trainable parameters in a module.""" |
|
return sum(p.numel() for p in module.parameters()) |
|
|
|
|
|
def download_file(path, url, digest=None): |
|
"""Downloads a file if it does not exist, optionally checking its SHA-256 hash.""" |
|
path = Path(path) |
|
path.parent.mkdir(parents=True, exist_ok=True) |
|
if not path.exists(): |
|
with urllib.request.urlopen(url) as response, open(path, 'wb') as f: |
|
shutil.copyfileobj(response, f) |
|
if digest is not None: |
|
file_digest = hashlib.sha256(open(path, 'rb').read()).hexdigest() |
|
if digest != file_digest: |
|
raise OSError(f'hash of {path} (url: {url}) failed to validate') |
|
return path |
|
|
|
|
|
@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. |
|
min_lr (float): The minimum 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., min_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.min_lr = min_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.min_lr, base_lr * lr_mult) |
|
for base_lr in self.base_lrs] |
|
|
|
|
|
class ExponentialLR(optim.lr_scheduler._LRScheduler): |
|
"""Implements an exponential learning rate schedule with an optional exponential |
|
warmup. When last_epoch=-1, sets initial lr as lr. Decays the learning rate |
|
continuously by decay (default 0.5) every num_steps steps. |
|
Args: |
|
optimizer (Optimizer): Wrapped optimizer. |
|
num_steps (float): The number of steps to decay the learning rate by decay in. |
|
decay (float): The factor by which to decay the learning rate every num_steps |
|
steps. Default: 0.5. |
|
warmup (float): Exponential warmup factor (0 <= warmup < 1, 0 to disable) |
|
Default: 0. |
|
min_lr (float): The minimum 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, num_steps, decay=0.5, warmup=0., min_lr=0., |
|
last_epoch=-1, verbose=False): |
|
self.num_steps = num_steps |
|
self.decay = decay |
|
if not 0. <= warmup < 1: |
|
raise ValueError('Invalid value for warmup') |
|
self.warmup = warmup |
|
self.min_lr = min_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 = (self.decay ** (1 / self.num_steps)) ** self.last_epoch |
|
return [warmup * max(self.min_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() |
|
|
|
|
|
def rand_v_diffusion(shape, sigma_data=1., min_value=0., max_value=float('inf'), device='cpu', dtype=torch.float32): |
|
"""Draws samples from a truncated v-diffusion training timestep distribution.""" |
|
min_cdf = math.atan(min_value / sigma_data) * 2 / math.pi |
|
max_cdf = math.atan(max_value / sigma_data) * 2 / math.pi |
|
u = torch.rand(shape, device=device, dtype=dtype) * (max_cdf - min_cdf) + min_cdf |
|
return torch.tan(u * math.pi / 2) * sigma_data |
|
|
|
|
|
def rand_split_log_normal(shape, loc, scale_1, scale_2, device='cpu', dtype=torch.float32): |
|
"""Draws samples from a split lognormal distribution.""" |
|
n = torch.randn(shape, device=device, dtype=dtype).abs() |
|
u = torch.rand(shape, device=device, dtype=dtype) |
|
n_left = n * -scale_1 + loc |
|
n_right = n * scale_2 + loc |
|
ratio = scale_1 / (scale_1 + scale_2) |
|
return torch.where(u < ratio, n_left, n_right).exp() |
|
|
|
|
|
class FolderOfImages(data.Dataset): |
|
"""Recursively finds all images in a directory. It does not support |
|
classes/targets.""" |
|
|
|
IMG_EXTENSIONS = {'.jpg', '.jpeg', '.png', '.ppm', '.bmp', '.pgm', '.tif', '.tiff', '.webp'} |
|
|
|
def __init__(self, root, transform=None): |
|
super().__init__() |
|
self.root = Path(root) |
|
self.transform = nn.Identity() if transform is None else transform |
|
self.paths = sorted(path for path in self.root.rglob('*') if path.suffix.lower() in self.IMG_EXTENSIONS) |
|
|
|
def __repr__(self): |
|
return f'FolderOfImages(root="{self.root}", len: {len(self)})' |
|
|
|
def __len__(self): |
|
return len(self.paths) |
|
|
|
def __getitem__(self, key): |
|
path = self.paths[key] |
|
with open(path, 'rb') as f: |
|
image = Image.open(f).convert('RGB') |
|
image = self.transform(image) |
|
return image, |
|
|
|
|
|
class CSVLogger: |
|
def __init__(self, filename, columns): |
|
self.filename = Path(filename) |
|
self.columns = columns |
|
if self.filename.exists(): |
|
self.file = open(self.filename, 'a') |
|
else: |
|
self.file = open(self.filename, 'w') |
|
self.write(*self.columns) |
|
|
|
def write(self, *args): |
|
print(*args, sep=',', file=self.file, flush=True) |
|
|
|
|
|
@contextmanager |
|
def tf32_mode(cudnn=None, matmul=None): |
|
"""A context manager that sets whether TF32 is allowed on cuDNN or matmul.""" |
|
cudnn_old = torch.backends.cudnn.allow_tf32 |
|
matmul_old = torch.backends.cuda.matmul.allow_tf32 |
|
try: |
|
if cudnn is not None: |
|
torch.backends.cudnn.allow_tf32 = cudnn |
|
if matmul is not None: |
|
torch.backends.cuda.matmul.allow_tf32 = matmul |
|
yield |
|
finally: |
|
if cudnn is not None: |
|
torch.backends.cudnn.allow_tf32 = cudnn_old |
|
if matmul is not None: |
|
torch.backends.cuda.matmul.allow_tf32 = matmul_old |
|
|
