PAIR/PAIR-Diffusion
38
1import importlib2 3import torch4from torch import optim5import numpy as np6 7from inspect import isfunction8from PIL import Image, ImageDraw, ImageFont9 10 11def log_txt_as_img(wh, xc, size=10):12 # wh a tuple of (width, height)13 # xc a list of captions to plot14 b = len(xc)15 txts = list()16 for bi in range(b):17 txt = Image.new("RGB", wh, color="white")18 draw = ImageDraw.Draw(txt)19 font = ImageFont.truetype('font/DejaVuSans.ttf', size=size)20 nc = int(40 * (wh[0] / 256))21 lines = "\n".join(xc[bi][start:start + nc] for start in range(0, len(xc[bi]), nc))22 23 try:24 draw.text((0, 0), lines, fill="black", font=font)25 except UnicodeEncodeError:26 print("Cant encode string for logging. Skipping.")27 28 txt = np.array(txt).transpose(2, 0, 1) / 127.5 - 1.029 txts.append(txt)30 txts = np.stack(txts)31 txts = torch.tensor(txts)32 return txts33 34 35def ismap(x):36 if not isinstance(x, torch.Tensor):37 return False38 return (len(x.shape) == 4) and (x.shape[1] > 3)39 40 41def isimage(x):42 if not isinstance(x,torch.Tensor):43 return False44 return (len(x.shape) == 4) and (x.shape[1] == 3 or x.shape[1] == 1)45 46 47def exists(x):48 return x is not None49 50 51def default(val, d):52 if exists(val):53 return val54 return d() if isfunction(d) else d55 56 57def mean_flat(tensor):58 """59 https://github.com/openai/guided-diffusion/blob/27c20a8fab9cb472df5d6bdd6c8d11c8f430b924/guided_diffusion/nn.py#L8660 Take the mean over all non-batch dimensions.61 """62 return tensor.mean(dim=list(range(1, len(tensor.shape))))63 64 65def count_params(model, verbose=False):66 total_params = sum(p.numel() for p in model.parameters())67 if verbose:68 print(f"{model.__class__.__name__} has {total_params*1.e-6:.2f} M params.")69 return total_params70 71 72def instantiate_from_config(config):73 if not "target" in config:74 if config == '__is_first_stage__':75 return None76 elif config == "__is_unconditional__":77 return None78 raise KeyError("Expected key `target` to instantiate.")79 return get_obj_from_str(config["target"])(**config.get("params", dict()))80 81 82def get_obj_from_str(string, reload=False):83 module, cls = string.rsplit(".", 1)84 if reload:85 module_imp = importlib.import_module(module)86 importlib.reload(module_imp)87 return getattr(importlib.import_module(module, package=None), cls)88 89 90class AdamWwithEMAandWings(optim.Optimizer):91 # credit to https://gist.github.com/crowsonkb/65f7265353f403714fce3b2595e0b29892 def __init__(self, params, lr=1.e-3, betas=(0.9, 0.999), eps=1.e-8, # TODO: check hyperparameters before using93 weight_decay=1.e-2, amsgrad=False, ema_decay=0.9999, # ema decay to match previous code94 ema_power=1., param_names=()):95 """AdamW that saves EMA versions of the parameters."""96 if not 0.0 <= lr:97 raise ValueError("Invalid learning rate: {}".format(lr))98 if not 0.0 <= eps:99 raise ValueError("Invalid epsilon value: {}".format(eps))100 if not 0.0 <= betas[0] < 1.0:101 raise ValueError("Invalid beta parameter at index 0: {}".format(betas[0]))102 if not 0.0 <= betas[1] < 1.0:103 raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1]))104 if not 0.0 <= weight_decay:105 raise ValueError("Invalid weight_decay value: {}".format(weight_decay))106 if not 0.0 <= ema_decay <= 1.0:107 raise ValueError("Invalid ema_decay value: {}".format(ema_decay))108 defaults = dict(lr=lr, betas=betas, eps=eps,109 weight_decay=weight_decay, amsgrad=amsgrad, ema_decay=ema_decay,110 ema_power=ema_power, param_names=param_names)111 super().__init__(params, defaults)112 113 def __setstate__(self, state):114 super().__setstate__(state)115 for group in self.param_groups:116 group.setdefault('amsgrad', False)117 118 @torch.no_grad()119 def step(self, closure=None):120 """Performs a single optimization step.121 Args:122 closure (callable, optional): A closure that reevaluates the model123 and returns the loss.124 """125 loss = None126 if closure is not None:127 with torch.enable_grad():128 loss = closure()129 130 for group in self.param_groups:131 params_with_grad = []132 grads = []133 exp_avgs = []134 exp_avg_sqs = []135 ema_params_with_grad = []136 state_sums = []137 max_exp_avg_sqs = []138 state_steps = []139 amsgrad = group['amsgrad']140 beta1, beta2 = group['betas']141 ema_decay = group['ema_decay']142 ema_power = group['ema_power']143 144 for p in group['params']:145 if p.grad is None:146 continue147 params_with_grad.append(p)148 if p.grad.is_sparse:149 raise RuntimeError('AdamW does not support sparse gradients')150 grads.append(p.grad)151 152 state = self.state[p]153 154 # State initialization155 if len(state) == 0:156 state['step'] = 0157 # Exponential moving average of gradient values158 state['exp_avg'] = torch.zeros_like(p, memory_format=torch.preserve_format)159 # Exponential moving average of squared gradient values160 state['exp_avg_sq'] = torch.zeros_like(p, memory_format=torch.preserve_format)161 if amsgrad:162 # Maintains max of all exp. moving avg. of sq. grad. values163 state['max_exp_avg_sq'] = torch.zeros_like(p, memory_format=torch.preserve_format)164 # Exponential moving average of parameter values165 state['param_exp_avg'] = p.detach().float().clone()166 167 exp_avgs.append(state['exp_avg'])168 exp_avg_sqs.append(state['exp_avg_sq'])169 ema_params_with_grad.append(state['param_exp_avg'])170 171 if amsgrad:172 max_exp_avg_sqs.append(state['max_exp_avg_sq'])173 174 # update the steps for each param group update175 state['step'] += 1176 # record the step after step update177 state_steps.append(state['step'])178 179 optim._functional.adamw(params_with_grad,180 grads,181 exp_avgs,182 exp_avg_sqs,183 max_exp_avg_sqs,184 state_steps,185 amsgrad=amsgrad,186 beta1=beta1,187 beta2=beta2,188 lr=group['lr'],189 weight_decay=group['weight_decay'],190 eps=group['eps'],191 maximize=False)192 193 cur_ema_decay = min(ema_decay, 1 - state['step'] ** -ema_power)194 for param, ema_param in zip(params_with_grad, ema_params_with_grad):195 ema_param.mul_(cur_ema_decay).add_(param.float(), alpha=1 - cur_ema_decay)196 197 return loss