Repository: sangyun884/blur-diffusion Branch: main Commit: 9dd8adb917bf Files: 27 Total size: 224.1 KB Directory structure: gitextract_8zm764cm/ ├── EMA.py ├── LICENSE ├── README.md ├── blur_diffusion.py ├── eval_x0hat.py ├── eval_x0hat.sh ├── fid.py ├── guided_diffusion/ │ ├── __init__.py │ ├── dist_util.py │ ├── fp16_util.py │ ├── gaussian_diffusion.py │ ├── image_datasets.py │ ├── logger.py │ ├── losses.py │ ├── nn.py │ ├── resample.py │ ├── respace.py │ ├── script_util.py │ ├── train_util.py │ └── unet.py ├── main.py ├── pytorch_fid/ │ ├── __init__.py │ ├── __main__.py │ ├── fid_score.py │ └── inception.py ├── train.sh └── utils.py ================================================ FILE CONTENTS ================================================ ================================================ FILE: EMA.py ================================================ # --------------------------------------------------------------- # Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. # # This work is licensed under the NVIDIA Source Code License # for Denoising Diffusion GAN. To view a copy of this license, see the LICENSE file. # --------------------------------------------------------------- ''' Codes adapted from https://github.com/NVlabs/LSGM/blob/main/util/ema.py ''' import warnings import torch from torch.optim import Optimizer class EMA(Optimizer): def __init__(self, opt, ema_decay): self.ema_decay = ema_decay self.apply_ema = self.ema_decay > 0. self.optimizer = opt self.state = opt.state self.param_groups = opt.param_groups def step(self, *args, **kwargs): retval = self.optimizer.step(*args, **kwargs) # stop here if we are not applying EMA if not self.apply_ema: return retval ema, params = {}, {} for group in self.optimizer.param_groups: for i, p in enumerate(group['params']): if p.grad is None: continue state = self.optimizer.state[p] # State initialization if 'ema' not in state: state['ema'] = p.data.clone() if p.shape not in params: params[p.shape] = {'idx': 0, 'data': []} ema[p.shape] = [] params[p.shape]['data'].append(p.data) ema[p.shape].append(state['ema']) for i in params: params[i]['data'] = torch.stack(params[i]['data'], dim=0) ema[i] = torch.stack(ema[i], dim=0) ema[i].mul_(self.ema_decay).add_(params[i]['data'], alpha=1. - self.ema_decay) for p in group['params']: if p.grad is None: continue idx = params[p.shape]['idx'] self.optimizer.state[p]['ema'] = ema[p.shape][idx, :] params[p.shape]['idx'] += 1 return retval def load_state_dict(self, state_dict): super(EMA, self).load_state_dict(state_dict) # load_state_dict loads the data to self.state and self.param_groups. We need to pass this data to # the underlying optimizer too. self.optimizer.state = self.state self.optimizer.param_groups = self.param_groups def swap_parameters_with_ema(self, store_params_in_ema): """ This function swaps parameters with their ema values. It records original parameters in the ema parameters, if store_params_in_ema is true.""" # stop here if we are not applying EMA if not self.apply_ema: warnings.warn('swap_parameters_with_ema was called when there is no EMA weights.') return for group in self.optimizer.param_groups: for i, p in enumerate(group['params']): if not p.requires_grad: continue ema = self.optimizer.state[p]['ema'] if store_params_in_ema: tmp = p.data.detach() p.data = ema.detach() self.optimizer.state[p]['ema'] = tmp else: p.data = ema.detach() ================================================ FILE: LICENSE ================================================ MIT License Copyright (c) 2022 Sangyun Lee Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. ================================================ FILE: README.md ================================================ # blur-diffusion This is the codebase for our paper Progressive Deblurring of Diffusion Models for Coarse-to-Fine Image Synthesis. ![Teaser image](./images/main_fig.jpg) > **Progressive Deblurring of Diffusion Models for Coarse-to-Fine Image Synthesis**
> Sangyun Lee1, Hyungjin Chung2, Jaehyeon Kim3, ‪Jong Chul Ye2 > 1Soongsil University, 2KAIST, 3Kakao Enterprise
> Paper: https://arxiv.org/abs/2207.11192
> **Abstract:** *Recently, diffusion models have shown remarkable results in image synthesis by gradually removing noise and amplifying signals. Although the simple generative process surprisingly works well, is this the best way to generate image data? For instance, despite the fact that human perception is more sensitive to the low-frequencies of an image, diffusion models themselves do not consider any relative importance of each frequency component. Therefore, to incorporate the inductive bias for image data, we propose a novel generative process that synthesizes images in a coarse-to-fine manner. First, we generalize the standard diffusion models by enabling diffusion in a rotated coordinate system with different velocities for each component of the vector. We further propose a blur diffusion as a special case, where each frequency component of an image is diffused at different speeds. Specifically, the proposed blur diffusion consists of a forward process that blurs an image and adds noise gradually, after which a corresponding reverse process deblurs an image and removes noise progressively. Experiments show that proposed model outperforms the previous method in FID on LSUN bedroom and church datasets.* ## Train ```bash train.sh``` ## Visualization ```bash eval_x0hat.sh``` ## Dataset Image files are required to compute FID during training. ================================================ FILE: blur_diffusion.py ================================================ import torch import torch.nn.functional as F import time import utils import os import numpy as np from utils import clear def gaussian_kernel_1d(kernel_size, sigma): assert sigma > 0.00001 x_cord = torch.arange(kernel_size) mean = (kernel_size - 1) / 2. # pdf of 1d gaussian not considering normalization constant gaussian_kernel = torch.exp(-0.5 * ((x_cord - mean)/sigma)**2) gaussian_kernel = gaussian_kernel / torch.sum(gaussian_kernel) return gaussian_kernel def betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar, max_beta=0.999): """ Ported from https://github.com/openai/guided-diffusion Utility function for cosine noise schedule """ betas = [] for i in range(num_diffusion_timesteps): t1 = i / num_diffusion_timesteps t2 = (i + 1) / num_diffusion_timesteps betas.append(min(1 - alpha_bar(t2) / alpha_bar(t1), max_beta)) return np.array(betas) class ExpSchedule: def __init__(self, N, offset = 1e-4): self.N = N def f(i): #return 1 - np.sin((i / N + 1) * np.pi / 2) + offset return np.exp(5*i/N - 5) + offset idxs = np.arange(N+1) self.alphas_bar = 1 - f(idxs) / f(idxs[-1]) self.alphas_bar_left_shifted = 1 - f(idxs-1) / f(idxs[-1]) self.alphas = self.alphas_bar / self.alphas_bar_left_shifted self.betas = 1 - self.alphas def get_betas(self): return self.betas # Increase blur strength as i increases class ForwardBlurIncreasing: def __init__(self, N, beta_min, beta_max, sig, sig_min, sig_max, D_diag, blur=None, noise_schedule='linear', channel=3, resolution=32, pad=None, device='cuda:0', f_type = 'linear'): # N: total number of discretizations used self.device = device self.N = N self.beta_min = beta_min self.beta_max = beta_max self.sig = sig self.sig_min = sig_min # \sigma_1 self.sig_max = sig_max # \sigma_N self.D_diag = D_diag # total dimension of image self.dim = self.D_diag.shape[0] # "deblur" class self.blur = blur self.resolution = resolution self.channel = channel self.pad = pad self.noise_schedule = noise_schedule if noise_schedule == 'linear': self.betas = torch.linspace(beta_min, beta_max, N, device=device) elif noise_schedule == 'cosine': self.betas = torch.tensor(betas_for_alpha_bar( N, lambda t: np.cos((t + 0.008) / 1.008 * np.pi / 2) ** 2, ), device=device, dtype=torch.float) elif noise_schedule == 'exp': Exp = ExpSchedule(N) betas = torch.tensor(Exp.get_betas()) self.betas = torch.tensor(betas, device=device) self.betas = F.pad(self.betas, (1, 0), value=0.0) self.alphas = 1 - self.betas self.sqrt_alphas = torch.sqrt(self.alphas) self.alphas_bar = torch.cumprod(self.alphas, dim=0) self.sqrt_alphas_bar = torch.cumprod(self.sqrt_alphas, dim=0) self.alphas_bar_prev = F.pad(self.alphas_bar, [1, 0], value=1)[:N] # def f(i): # a = (self.sig_max**2 - self.sig_min**2) / (self.sig**2 * (self.N**2 - 1)) # b = (self.sig_min / self.sig) ** 2 - a # return a * i**2 + b def f_linear(i): f_N = (self.sig_max / self.sig) ** 2 f_1 = (self.sig_min / self.sig) ** 2 return (f_N - f_1)/ (N-1) * (i - 1) + f_1 idxs = torch.tensor(np.arange(N+1), device=device, dtype=torch.float) def f_log(i): log = lambda x: torch.log(x + 1e-6) / (10*np.log(N)) f_N = (self.sig_max / self.sig) ** 2 f_1 = (self.sig_min / self.sig) ** 2 a = (f_N - f_1) / log(torch.tensor(N)) b = f_1 return a*log(i) + b def f_quadratic(i): f_N = (sig_max / sig) ** 2 f_1 = (sig_min / sig) ** 2 a = (f_N - f_1) / (N**2 - 1) b = f_1 - (f_N - f_1) / (N**2 - 1) return a*i**2 + b def cubic(i): f_N = (sig_max / sig) ** 2 f_1 = (sig_min / sig) ** 2 a = (f_N-f_1) / N**3 b = f_1 return a*i**3 + b def quartic(i): f_N = (sig_max / sig) ** 2 f_1 = (sig_min / sig) ** 2 a = (f_N-f_1) / N**4 b = f_1 return a*i**4 + b def triangular(i): less_than_N_2 = (i < N/2).type(torch.int) return f_linear(i) * less_than_N_2 + f_linear(N-i) * (1-less_than_N_2 ) if f_type == 'linear': f = f_linear elif f_type == 'log': f = f_log elif f_type == 'quadratic': f = f_quadratic elif f_type == 'cubic': f = cubic elif f_type == 'quartic': f = quartic elif f_type == 'triangular': f = triangular else: raise NotImplementedError self.fs = f(idxs) print("fs: ", self.fs) self.fs_cum = torch.cumsum(self.fs, dim=0) idxs_long = torch.tensor(np.arange(N+1), device=device, dtype=torch.long) def B(i): p = (2 * f(i)).unsqueeze(-1).repeat(1, self.D_diag.shape[0]) print("p: ", p.shape) D = self.D_diag.repeat(p.shape[0], 1) print("D: ", D.shape) return self.alphas[i].unsqueeze(-1) * self.D_diag ** p self.Bs = B(idxs_long) self.Bs_bar = F.pad(torch.cumprod(self.Bs[1:], dim=0), (0,0,1, 0), value=0) self.one_minus_Bs_bar = 1 - self.Bs_bar self.one_minus_Bs_bar_sqrt = torch.sqrt(self.one_minus_Bs_bar) self.Bs_sqrt = torch.sqrt(self.Bs) self.Bs_squared = self.Bs ** 2 self.Bs_bar_sqrt = torch.sqrt(self.Bs_bar) # fname = f'res{resolution}_N{N}_sig{sig}_sigmin_{sig_min}_sigmax_{sig_max}_b_mat_sq-nc{channel}_{noise_schedule}_inc_nopad_rotated.npz' # if os.path.isfile(fname): # with np.load(fname, allow_pickle=True) as data: # print("data", list(data.keys())) # data = data['arr'].astype(np.float32) # b_mat_sq = torch.tensor(data, device=device) # self.b_mat_sq = b_mat_sq # else: # tic = time.time() # self.b_mat_sq = torch.zeros([N + 1, self.dim]).to(device) # self.b_mat = torch.sqrt(self.b_mat_sq) def _beta_i(self, i): return self.betas[i] def _alpha_i(self, i): return self.alphas[i] def _alphas_bar_i(self, i): return self.alphas_bar[i] def _Bhat_i(self, i): return self.b_mat[i, :] # Forward blur def _Bhat_sq_i(self, i): assert i != 0 return self.b_mat_sq[i, :] # Forward blur def get_mean(self, x0, i): mat = self.Bs_bar_sqrt[i] mean = self.blur.U(mat * self.blur.Ut(x0)) return mean # Forward blur def get_std(self,i, noise): mat = self.one_minus_Bs_bar_sqrt[i] std = self.blur.U(mat * self.blur.Ut(noise)) return std # Forward blur def get_var(self, x0, i, noise=None): raise NotImplementedError Bhat_i = self._Bhat_sq_i(i) noise = noise if noise is not None else torch.randn_like(x0) std = self.blur.U(Bhat_i * self.blur.Ut(noise)) return std def W(self, x, i): mat = self.Bs_sqrt[i] batch = x.shape[0] blurred = self.blur.U(mat * self.blur.Ut(x)).view(batch, self.channel, self.resolution, self.resolution) return blurred def W_inv(self, x, i): mat = self.Bs_squared[i] batch = x.shape[0] blurred = self.blur.U(mat * self.blur.Ut(x)).view(batch, self.channel, self.resolution, self.resolution) return blurred def U_I_minus_B_Ut(self, x, i): # U(I-B)UT batch = x.shape[0] mat = 1 - self.Bs[i] result = self.blur.U(mat * self.blur.Ut(x)).view(batch, self.channel, self.resolution, self.resolution) return result def U_I_minus_B_sqrt_Ut(self, x, i): # U(sqrt(I-B))UT batch = x.shape[0] mat = torch.sqrt(1 - self.Bs[i]) result = self.blur.U(mat * self.blur.Ut(x)).view(batch, self.channel, self.resolution, self.resolution) return result # Forward blur def get_x_i(self, x0, i, return_eps = False): assert 0 not in i if len(x0.shape) == 4: batch = x0.shape[0] # reparam trick mean = self.get_mean(x0, i) noise = torch.randn_like(x0) std = self.get_std(i, noise = noise) # crop if len(x0.shape) == 4: img = (mean + std).view(batch, self.channel, self.resolution, self.resolution) else: img = (mean + std).view(self.channel, self.resolution, self.resolution) # cropped = cropped + torch.randn_like(cropped) * torch.sqrt(self.betas[i-1]) if return_eps: return img, noise return img def get_x_N(self, x0_shape, N): """ prior sampling """ if len(x0_shape) == 4: batch = x0_shape[0] # pad x0 = torch.zeros(x0_shape, device=self.device) noise = torch.randn_like(x0) std = self.get_std(N, noise) # crop if len(x0_shape) == 4: img = std.view(batch, self.channel, self.resolution, self.resolution) else: img = std.view(self.channel, self.resolution, self.resolution) return img def get_x0_from_eps(self, xi, eps, i): batch = xi.shape[0] std = self.get_std(i, noise = eps).view(batch, self.channel, self.resolution, self.resolution) mean = xi - std return mean / torch.sqrt(self.alphas_bar[i]).view(-1, 1, 1, 1) mat = self.Bs_bar_sqrt[i] ** (-1) x0 = self.blur.U(mat * self.blur.Ut(mean)).view(batch, self.channel, self.resolution, self.resolution) return x0 def get_score_gt(self, xi, x0, i): # Return the ground-truth score of x_i assert torch.all(i>0) batch = x0.shape[0] mean = self.get_mean(x0, i).view(batch, self.channel, self.resolution, self.resolution) diff = xi - mean mat = (self.one_minus_Bs_bar[i])**(-1) score = -self.blur.U(mat * self.blur.Ut(diff)).view(batch, self.channel, self.resolution, self.resolution) return score def get_score_gt2(self, xi, x0, i): # Return the ground-truth score of x_i batch = x0.shape[0] x0_pad = utils.pad(x0, self.pad) # a_iW^ix_0 mean_pad = self.get_mean(x0_pad, i).view(batch, self.channel, self.resolution + self.pad * 2, self.resolution + self.pad * 2) mean = utils.crop(mean_pad, self.pad) diff = xi - mean diff = utils.pad(diff, self.pad) score = self.blur.U(self._Bhat_i(i) ** (-2) * self.blur.Ut(diff)).view(batch, self.channel, self.resolution + self.pad * 2, self.resolution + self.pad * 2) score = -utils.crop(score, self.pad) return score def sanity(self, x_0, i): batch = x_0.shape[0] tol = 1e-2 xi, eps = self.get_x_i(x_0, i, return_eps = True) score1 = self.get_score_from_eps(eps, i) score2 = self.get_score_gt(xi, x_0, i) # utils.tensor_imsave(score1[0], "./", "score1.jpg") # utils.tensor_imsave(score2[0], "./", "score2.jpg") # rms = lambda x: torch.sqrt(torch.mean(x**2)) # print(f"rms(score1) = {rms(score1)}, rms(score2) = {rms(score2)}") # noise = torch.randn_like(x_0) # noise_pad = utils.pad(noise, self.pad) # std = self.blur.U(self._Bhat_i(i) * self.blur.Ut(noise_pad)).view(batch, # self.channel, # self.resolution + self.pad * 2, # self.resolution + self.pad * 2) # std = utils.crop(std, self.pad) # std = utils.pad(std, self.pad) # # crop -> pad makes error # score1 = self.blur.U(self._Bhat_i(i) ** (-2) * self.blur.Ut(std)).view(batch, # self.channel, # self.resolution + self.pad * 2, # self.resolution + self.pad * 2) # score1 = -utils.crop(score1, self.pad) # score2 = self.blur.U(self._Bhat_i(i) ** (-1) * self.blur.Ut(noise_pad)).view(batch, self.channel, self.resolution + self.pad * 2, self.resolution + self.pad * 2) # score2 = -utils.crop(score2, self.pad) MAE = torch.mean(torch.abs(score1 - score2)) assert MAE < tol, f'MAE = {MAE}' print(f"MAE = {MAE}") # score3 = self.get_score_gt2(xi, x_0, i) # print(f"score3 = {score3[:5]}") # print(f"score2 = {score2[:5]}") # MAE = torch.mean(torch.abs(score3 - score2)) # print(f"MAE = {MAE}, fraction = {MAE/torch.mean(torch.abs(score3))}, norm: {torch.mean(torch.abs(score3))}") # assert MAE <1e-5, f'MAE = {MAE}' # MAE = torch.mean(torch.abs(self.b_mat_sq[1:] - self.one_minus_Bs_bar[1:])) # assert MAE <1e-5, f'MAE = {MAE}' def get_score_from_eps(self, eps, i): # Return the score of x_i batch = eps.shape[0] mat = self.one_minus_Bs_bar_sqrt[i]**(-1) score = -self.blur.U(mat * self.blur.Ut(eps)).view(batch, self.channel, self.resolution, self.resolution) return score def get_score_from_std(self, std, i): # Return the score of x_i batch = std.shape[0] mat = self.one_minus_Bs_bar[i] ** (-1) score = -self.blur.U(mat * self.blur.Ut(std)).view(batch, self.channel, self.resolution, self.resolution) return score def get_loss_i_exact(self, model, x0, xi, i): # Calculate the DSM objective (exact) batch = x0.shape[0] pred = model(xi, i) score = self.get_score_gt(x0, xi, i) loss = torch.mean((pred - score) ** 2) return loss def get_loss_i_simple(self, model, x0, xi, i): raise NotImplementedError # Calculate the DSM objective (simple) batch = x0.shape[0] pred = model(xi, i) # UB**2U^Ts pred = utils.pad(pred, self.pad) pred = self.blur.U(self._Bhat_i(i) ** 2 * self.blur.Ut(pred)).view(batch, self.channel, self.resolution + self.pad * 2, self.resolution + self.pad * 2) pred = utils.crop(pred, self.pad) # ai[UDUT]x0 x0_pad = utils.pad(x0, self.pad) mean = self.get_mean(x0_pad, i).view(batch, self.channel, self.resolution + self.pad * 2, self.resolution + self.pad * 2) mean = utils.crop(mean, self.pad) loss = torch.mean((pred + xi - mean) ** 2) return loss def get_loss_i_eps_simple(self, model, x_i, i, eps): pred = model(x_i, i) loss = torch.mean((pred - eps) ** 2) return loss def get_loss_i_std_matching(self, model, x_i, i, eps): pred = model(x_i, i) batch = x_i.shape[0] std = self.get_std(i, noise = eps).view(batch, self.channel, self.resolution, self.resolution) loss = torch.mean((pred - std) ** 2) return loss class H_functions: """ Ported from https://github.com/bahjat-kawar/ddrm A class replacing the SVD of a matrix H, perhaps efficiently. All input vectors are of shape (Batch, ...). All output vectors are of shape (Batch, DataDimension). """ def V(self, vec): """ Multiplies the input vector by V """ raise NotImplementedError() def Vt(self, vec): """ Multiplies the input vector by V transposed """ raise NotImplementedError() def U(self, vec): """ Multiplies the input vector by U """ raise NotImplementedError() def Ut(self, vec): """ Multiplies the input vector by U transposed """ raise NotImplementedError() def singulars(self): """ Returns a vector containing the singular values. The shape of the vector should be the same as the smaller dimension (like U) """ raise NotImplementedError() def add_zeros(self, vec): """ Adds trailing zeros to turn a vector from the small dimension (U) to the big dimension (V) """ raise NotImplementedError() def H(self, vec): """ Multiplies the input vector by H """ temp = self.Vt(vec) singulars = self.singulars() return self.U(singulars * temp[:, :singulars.shape[0]]) def Ht(self, vec): """ Multiplies the input vector by H transposed """ temp = self.Ut(vec) singulars = self.singulars() return self.V(self.add_zeros(singulars * temp[:, :singulars.shape[0]])) def H_pinv(self, vec): """ Multiplies the input vector by the pseudo inverse of H """ temp = self.Ut(vec) singulars = self.singulars() temp[:, :singulars.shape[0]] = temp[:, :singulars.shape[0]] / singulars return self.V(self.add_zeros(temp)) class Deblurring(H_functions): def mat_by_img(self, M, v): return torch.matmul(M, v.reshape(v.shape[0] * self.channels, self.img_dim, self.img_dim)).reshape(v.shape[0], self.channels, M.shape[0], self.img_dim) def img_by_mat(self, v, M): return torch.matmul(v.reshape(v.shape[0] * self.channels, self.img_dim, self.img_dim), M).reshape(v.shape[0], self.channels, self.img_dim, M.shape[1]) def __init__(self, kernel, channels, img_dim, device): self.img_dim = img_dim self.channels = channels # build 1D conv matrix H_small = torch.zeros(img_dim, img_dim, device=device) for i in range(img_dim): for j in range(i - kernel.shape[0] // 2, i + kernel.shape[0] // 2): if j < 0 or j >= img_dim: continue H_small[i, j] = kernel[j - i + kernel.shape[0] // 2] # torch.cholesky(H_small) # raise error if H_small is not positive definite self.H_small = H_small # positive definite matrix # get the evd of the 1D conv self.U_small, self.singulars_small, _ = torch.svd(H_small, some=False) self.V_small = self.U_small # V should equal U since H is symmetric ZERO = 3e-2 # truncate self.singulars_small[self.singulars_small < ZERO] = ZERO # calculate the singular values of the big matrix using Kronecker product self._singulars = torch.matmul(self.singulars_small.reshape(img_dim, 1), self.singulars_small.reshape(1, img_dim)).reshape(img_dim ** 2) self._singulars[self._singulars > 1] = 1 print("contains zero?", torch.any(self._singulars == 0)) # raise NotImplementedError # sort the big matrix singulars and save the permutation self._singulars, self._perm = self._singulars.sort(descending=True) # , stable=True) def V(self, vec): # invert the permutation temp = torch.zeros(vec.shape[0], self.img_dim ** 2, self.channels, device=vec.device) temp[:, self._perm, :] = vec.clone().reshape(vec.shape[0], self.img_dim ** 2, self.channels) temp = temp.permute(0, 2, 1) # multiply the image by V from the left and by V^T from the right out = self.mat_by_img(self.V_small, temp) out = self.img_by_mat(out, self.V_small.transpose(0, 1)).reshape(vec.shape[0], -1) return out def Vt(self, vec): # multiply the image by V^T from the left and by V from the right temp = self.mat_by_img(self.V_small.transpose(0, 1), vec.clone()) temp = self.img_by_mat(temp, self.V_small).reshape(vec.shape[0], self.channels, -1) # permute the entries according to the singular values temp = temp[:, :, self._perm].permute(0, 2, 1) return temp.reshape(vec.shape[0], -1) def U(self, vec): # invert the permutation temp = torch.zeros(vec.shape[0], self.img_dim ** 2, self.channels, device=vec.device) temp[:, self._perm, :] = vec.clone().reshape(vec.shape[0], self.img_dim ** 2, self.channels) temp = temp.permute(0, 2, 1) # multiply the image by U from the left and by U^T from the right out = self.mat_by_img(self.U_small, temp) out = self.img_by_mat(out, self.U_small.transpose(0, 1)).reshape(vec.shape[0], -1) return out def Ut(self, vec): # multiply the image by U^T from the left and by U from the right temp = self.mat_by_img(self.U_small.transpose(0, 1), vec.clone()) temp = self.img_by_mat(temp, self.U_small).reshape(vec.shape[0], self.channels, -1) # permute the entries according to the singular values temp = temp[:, :, self._perm].permute(0, 2, 1) return temp.reshape(vec.shape[0], -1) def conv1d_col_matmul(self, x): return torch.matmul(self.H_small, x) def conv1d_row_matmul(self, x): return torch.matmul(x, self.H_small) def conv2d_sep_matmul(self, x): return torch.matmul(torch.matmul(self.H_small, x), self.H_small) def singulars(self): if self.channels == 3: return self._singulars.repeat(1, 3).reshape(-1) else: return self._singulars.reshape(-1) def update_singulars(self, new_singulars): self._singulars = new_singulars def add_zeros(self, vec): return vec.clone().reshape(vec.shape[0], -1) ================================================ FILE: eval_x0hat.py ================================================ import torch from PIL import Image from collections import defaultdict import torch.nn.functional as TF import torchvision.datasets as dsets from torchvision import transforms import numpy as np import torch.optim as optim import torchvision import matplotlib.pyplot as plt import utils from guided_diffusion.unet import UNetModel import math from tensorboardX import SummaryWriter import os import json from collections import namedtuple import argparse from torchvision.utils import save_image from tqdm import tqdm from blur_diffusion import Deblurring, ForwardBlurIncreasing, gaussian_kernel_1d from utils import normalize_np, clear from EMA import EMA from torch.nn import DataParallel parser = argparse.ArgumentParser(description='Configs') parser.add_argument('--gpu', type=str, help='gpu num') parser.add_argument('--name', type=str, help='Saving directory name') parser.add_argument('--ckpt', default='', type=str, help='UNet checkpoint') parser.add_argument('--bsize', default=16, type=int, help='batchsize') parser.add_argument('--N', default=1000, type=int, help='Max diffusion timesteps') parser.add_argument('--sig', default=0.3, type=float, help='sigma value for blur kernel') parser.add_argument('--sig_min', default=0.3, type=float, help='sigma value for blur kernel') parser.add_argument('--sig_max', default=1.5, type=float, help='sigma value for blur kernel') parser.add_argument('--noise_schedule', default='linear', type=str, help='Type of noise schedule to use') parser.add_argument('--betamin', default=0.0001, type=float, help='beta (min). get_score(1) can diverge if this is too low.') parser.add_argument('--betamax', default=0.02, type=float, help='beta (max)') parser.add_argument('--res', type=int, help='resolution') parser.add_argument('--nc', type=int, help='channels') parser.add_argument('--loss_type', type=str, default = 'sm_simple', choices=['sm_simple', 'eps_simple', 'sm_exact']) parser.add_argument('--f_type', type=str, default = 'linear', choices=['linear', 'log', 'quadratic', 'quartic']) parser.add_argument('--dropout', default=0, type=float, help='dropout') parser.add_argument('--freq_feat', action='store_true', help = "concat Utx_i") opt = parser.parse_args() ksize = opt.res * 2 - 1 pad = 0 bsize = opt.bsize beta_min = opt.betamin beta_max = opt.betamax device = torch.device(f'cuda:{opt.gpu}') # define forward blur kernel = gaussian_kernel_1d(ksize, opt.sig) blur = Deblurring(kernel, opt.nc, opt.res, device=device) print("blur.U_small.shape:", blur.U_small.shape) D_diag = blur.singulars() fb = ForwardBlurIncreasing(N=opt.N, beta_min=beta_min, beta_max=beta_max, sig=opt.sig, sig_max = opt.sig_max, sig_min = opt.sig_min, D_diag=D_diag, blur=blur, channel=opt.nc, device=device, noise_schedule=opt.noise_schedule, resolution=opt.res, pad=pad, f_type=opt.f_type) iter = opt.ckpt.split('/')[-1].split('.')[0] dir = os.path.join('experiments', opt.name, f'inferences-{iter}') if not os.path.exists(dir): os.mkdir(dir) model = UNetModel(opt.res, opt.nc, 128, opt.nc, blur = blur, dropout=opt.dropout, freq_feat = opt.freq_feat) model.load_state_dict(torch.load(opt.ckpt)) model.to(device) model.eval() print("input_nc", opt.nc, "resolution", opt.res) num_samples = bsize id = 0 for _ in tqdm(range(num_samples // bsize)): with torch.no_grad(): i = np.array([opt.N - 1] * bsize) i = torch.from_numpy(i).to(device) pred = fb.get_x_N([bsize, opt.nc, opt.res, opt.res], i) x0_list = [] for i in reversed(range(1, opt.N)): i = np.array([i] * bsize) i = torch.from_numpy(i).to(device) if opt.loss_type == "sm_simple": s = model(pred, i) elif opt.loss_type == "eps_simple": eps = model(pred, i) s = fb.get_score_from_eps(eps, i) x0_hat = fb.get_x0_from_eps(pred, eps, i) if i[0] % 100 == 0 or i[0] == 1: x0_list.append(x0_hat) elif opt.loss_type == "sm_exact": s = model(pred, i) else: raise NotImplementedError s = fb.U_I_minus_B_Ut(s, i) rms = lambda x: torch.sqrt(torch.mean(x ** 2)) print(f"rms(s) * fb._beta_i(i) = {rms(s) * fb._beta_i(i)[0]}") hf = pred - fb.W(pred, i) # Anderson theorem pred1 = pred + hf # unsharpening mask filtering pred2 = pred1 + s # # denoising if i[0] > 2: pred = pred2 + fb.U_I_minus_B_sqrt_Ut(torch.randn_like(pred), i) # inject noise else: pred = pred2 print(f"i = {i[0]}, rmse = {torch.sqrt(torch.mean(pred**2))}, mean = {torch.mean(pred)} std = {torch.std(pred)}" ) grid = torch.cat(x0_list, dim=3) # grid_sample.shape: (bsize, channel, H, W * 12) # (batch_size, channel, H, W * 12) -> (channel, H * bsize, W * 12) grid = grid.permute(1, 0, 2, 3).contiguous().view(grid.shape[1], -1, grid.shape[3]) # (bsize, channel, H, W) -> (channel, H, W * bsize) save_image(grid, os.path.join(dir, f'xhat.png')) ================================================ FILE: eval_x0hat.sh ================================================ #!/bin/bash export OMP_NUM_THREADS=1 export CUDA_VISIBLE_DEVICES=1 python eval_x0hat.py \ --name test \ --noise_schedule linear \ --gpu 0 \ --bsize 4 \ --sig 0.4 \ --sig_min 0 \ --sig_max 0.15 \ --f_type quartic \ --res 64 \ --nc 3 \ --loss_type eps_simple \ --ckpt /home/ubuntu/code/sangyoon/forward-blur-new/experiments/lsun-simplified-64-sigmax0.15-quartic/model_450001.ckpt ================================================ FILE: fid.py ================================================ from pytorch_fid.fid_score import compute_statistics_of_path, calculate_frechet_distance from pytorch_fid.inception import InceptionV3 import os class FID(object): def __init__(self, real_dir, device, bsize = 128): # https://github.com/mseitzer/pytorch-fid # If real_dir contains a .npz file, then we'll just use that. # Otherwise it's assumed to be a directory. # In that case, statistics will be computed, and .npz file will be saved. if os.path.exists(os.path.join(real_dir, "inception_statistics.npz")): real_dir = os.path.join(real_dir, "inception_statistics.npz") self.dims = 2048 self.bsize = bsize self.device = device block_idx = InceptionV3.BLOCK_INDEX_BY_DIM[self.dims] self.model = InceptionV3([block_idx]).to(device) m, s = compute_statistics_of_path(real_dir, self.model, bsize, self.dims, device) self.mu_real = m self.sig_real = s def __call__(self, fake_dir): m, s = compute_statistics_of_path(fake_dir, self.model, self.bsize, self.dims, self.device) fid_value = calculate_frechet_distance(m, s, self.mu_real, self.sig_real) return fid_value ================================================ FILE: guided_diffusion/__init__.py ================================================ """ Codebase for "Improved Denoising Diffusion Probabilistic Models". """ ================================================ FILE: guided_diffusion/dist_util.py ================================================ """ Helpers for distributed training. """ import io import os import socket import blobfile as bf from mpi4py import MPI import torch as th import torch.distributed as dist # Change this to reflect your cluster layout. # The GPU for a given rank is (rank % GPUS_PER_NODE). GPUS_PER_NODE = 8 SETUP_RETRY_COUNT = 3 def setup_dist(): """ Setup a distributed process group. """ if dist.is_initialized(): return os.environ["CUDA_VISIBLE_DEVICES"] = f"{MPI.COMM_WORLD.Get_rank() % GPUS_PER_NODE}" comm = MPI.COMM_WORLD backend = "gloo" if not th.cuda.is_available() else "nccl" # backend = "nccl" if backend == "gloo": hostname = "localhost" else: hostname = socket.gethostbyname(socket.getfqdn()) # print(hostname) os.environ["MASTER_ADDR"] = comm.bcast(hostname, root=0) os.environ["RANK"] = str(comm.rank) os.environ["WORLD_SIZE"] = str(comm.size) port = comm.bcast(_find_free_port(), root=0) os.environ["MASTER_PORT"] = str(port) # print(port) # print(os.environ["MASTER_ADDR"], os.environ["RANK"], os.environ["WORLD_SIZE"]) print(r"dist.init_process_group(backend=backend, init_method=\"env://\") -- start") dist.init_process_group(backend=backend, init_method="env://") print(r"dist.init_process_group(backend=backend, init_method=\"env://\") -- end") def dev(): """ Get the device to use for torch.distributed. """ if th.cuda.is_available(): return th.device(f"cuda") return th.device("cpu") def load_state_dict(path, **kwargs): """ Load a PyTorch file without redundant fetches across MPI ranks. """ chunk_size = 2 ** 30 # MPI has a relatively small size limit if MPI.COMM_WORLD.Get_rank() == 0: with bf.BlobFile(path, "rb") as f: data = f.read() num_chunks = len(data) // chunk_size if len(data) % chunk_size: num_chunks += 1 MPI.COMM_WORLD.bcast(num_chunks) for i in range(0, len(data), chunk_size): MPI.COMM_WORLD.bcast(data[i : i + chunk_size]) else: num_chunks = MPI.COMM_WORLD.bcast(None) data = bytes() for _ in range(num_chunks): data += MPI.COMM_WORLD.bcast(None) return th.load(io.BytesIO(data), **kwargs) def sync_params(params): """ Synchronize a sequence of Tensors across ranks from rank 0. """ for p in params: with th.no_grad(): dist.broadcast(p, 0) def _find_free_port(): try: s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) s.bind(("", 0)) s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) return s.getsockname()[1] finally: s.close() ================================================ FILE: guided_diffusion/fp16_util.py ================================================ """ Helpers to train with 16-bit precision. """ import numpy as np import torch as th import torch.nn as nn from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors from . import logger INITIAL_LOG_LOSS_SCALE = 20.0 def convert_module_to_f16(l): """ Convert primitive modules to float16. """ if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Conv3d)): l.weight.data = l.weight.data.half() if l.bias is not None: l.bias.data = l.bias.data.half() def convert_module_to_f32(l): """ Convert primitive modules to float32, undoing convert_module_to_f16(). """ if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Conv3d)): l.weight.data = l.weight.data.float() if l.bias is not None: l.bias.data = l.bias.data.float() def make_master_params(param_groups_and_shapes): """ Copy model parameters into a (differently-shaped) list of full-precision parameters. """ master_params = [] for param_group, shape in param_groups_and_shapes: master_param = nn.Parameter( _flatten_dense_tensors( [param.detach().float() for (_, param) in param_group] ).view(shape) ) master_param.requires_grad = True master_params.append(master_param) return master_params def model_grads_to_master_grads(param_groups_and_shapes, master_params): """ Copy the gradients from the model parameters into the master parameters from make_master_params(). """ for master_param, (param_group, shape) in zip( master_params, param_groups_and_shapes ): master_param.grad = _flatten_dense_tensors( [param_grad_or_zeros(param) for (_, param) in param_group] ).view(shape) def master_params_to_model_params(param_groups_and_shapes, master_params): """ Copy the master parameter data back into the model parameters. """ # Without copying to a list, if a generator is passed, this will # silently not copy any parameters. for master_param, (param_group, _) in zip(master_params, param_groups_and_shapes): for (_, param), unflat_master_param in zip( param_group, unflatten_master_params(param_group, master_param.view(-1)) ): param.detach().copy_(unflat_master_param) def unflatten_master_params(param_group, master_param): return _unflatten_dense_tensors(master_param, [param for (_, param) in param_group]) def get_param_groups_and_shapes(named_model_params): named_model_params = list(named_model_params) scalar_vector_named_params = ( [(n, p) for (n, p) in named_model_params if p.ndim <= 1], (-1), ) matrix_named_params = ( [(n, p) for (n, p) in named_model_params if p.ndim > 1], (1, -1), ) return [scalar_vector_named_params, matrix_named_params] def master_params_to_state_dict( model, param_groups_and_shapes, master_params, use_fp16 ): if use_fp16: state_dict = model.state_dict() for master_param, (param_group, _) in zip( master_params, param_groups_and_shapes ): for (name, _), unflat_master_param in zip( param_group, unflatten_master_params(param_group, master_param.view(-1)) ): assert name in state_dict state_dict[name] = unflat_master_param else: state_dict = model.state_dict() for i, (name, _value) in enumerate(model.named_parameters()): assert name in state_dict state_dict[name] = master_params[i] return state_dict def state_dict_to_master_params(model, state_dict, use_fp16): if use_fp16: named_model_params = [ (name, state_dict[name]) for name, _ in model.named_parameters() ] param_groups_and_shapes = get_param_groups_and_shapes(named_model_params) master_params = make_master_params(param_groups_and_shapes) else: master_params = [state_dict[name] for name, _ in model.named_parameters()] return master_params def zero_master_grads(master_params): for param in master_params: param.grad = None def zero_grad(model_params): for param in model_params: # Taken from https://pytorch.org/docs/stable/_modules/torch/optim/optimizer.html#Optimizer.add_param_group if param.grad is not None: param.grad.detach_() param.grad.zero_() def param_grad_or_zeros(param): if param.grad is not None: return param.grad.data.detach() else: return th.zeros_like(param) class MixedPrecisionTrainer: def __init__( self, *, model, use_fp16=False, fp16_scale_growth=1e-3, initial_lg_loss_scale=INITIAL_LOG_LOSS_SCALE, ): self.model = model self.use_fp16 = use_fp16 self.fp16_scale_growth = fp16_scale_growth self.model_params = list(self.model.parameters()) self.master_params = self.model_params self.param_groups_and_shapes = None self.lg_loss_scale = initial_lg_loss_scale if self.use_fp16: self.param_groups_and_shapes = get_param_groups_and_shapes( self.model.named_parameters() ) self.master_params = make_master_params(self.param_groups_and_shapes) self.model.convert_to_fp16() def zero_grad(self): zero_grad(self.model_params) def backward(self, loss: th.Tensor): if self.use_fp16: loss_scale = 2 ** self.lg_loss_scale (loss * loss_scale).backward() else: loss.backward() def optimize(self, opt: th.optim.Optimizer): if self.use_fp16: return self._optimize_fp16(opt) else: return self._optimize_normal(opt) def _optimize_fp16(self, opt: th.optim.Optimizer): logger.logkv_mean("lg_loss_scale", self.lg_loss_scale) model_grads_to_master_grads(self.param_groups_and_shapes, self.master_params) grad_norm, param_norm = self._compute_norms(grad_scale=2 ** self.lg_loss_scale) if check_overflow(grad_norm): self.lg_loss_scale -= 1 logger.log(f"Found NaN, decreased lg_loss_scale to {self.lg_loss_scale}") zero_master_grads(self.master_params) return False logger.logkv_mean("grad_norm", grad_norm) logger.logkv_mean("param_norm", param_norm) self.master_params[0].grad.mul_(1.0 / (2 ** self.lg_loss_scale)) opt.step() zero_master_grads(self.master_params) master_params_to_model_params(self.param_groups_and_shapes, self.master_params) self.lg_loss_scale += self.fp16_scale_growth return True def _optimize_normal(self, opt: th.optim.Optimizer): grad_norm, param_norm = self._compute_norms() logger.logkv_mean("grad_norm", grad_norm) logger.logkv_mean("param_norm", param_norm) opt.step() return True def _compute_norms(self, grad_scale=1.0): grad_norm = 0.0 param_norm = 0.0 for p in self.master_params: with th.no_grad(): param_norm += th.norm(p, p=2, dtype=th.float32).item() ** 2 if p.grad is not None: grad_norm += th.norm(p.grad, p=2, dtype=th.float32).item() ** 2 return np.sqrt(grad_norm) / grad_scale, np.sqrt(param_norm) def master_params_to_state_dict(self, master_params): return master_params_to_state_dict( self.model, self.param_groups_and_shapes, master_params, self.use_fp16 ) def state_dict_to_master_params(self, state_dict): return state_dict_to_master_params(self.model, state_dict, self.use_fp16) def check_overflow(value): return (value == float("inf")) or (value == -float("inf")) or (value != value) ================================================ FILE: guided_diffusion/gaussian_diffusion.py ================================================ """ This code started out as a PyTorch port of Ho et al's diffusion models: https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/diffusion_utils_2.py Docstrings have been added, as well as DDIM sampling and a new collection of beta schedules. """ import enum import math import numpy as np import torch as th from torchvision import transforms from .nn import mean_flat from .losses import normal_kl, discretized_gaussian_log_likelihood import os from PIL import Image def get_named_beta_schedule(schedule_name, num_diffusion_timesteps): """ Get a pre-defined beta schedule for the given name. The beta schedule library consists of beta schedules which remain similar in the limit of num_diffusion_timesteps. Beta schedules may be added, but should not be removed or changed once they are committed to maintain backwards compatibility. """ if schedule_name == "linear": # Linear schedule from Ho et al, extended to work for any number of # diffusion steps. scale = 1000 / num_diffusion_timesteps beta_start = scale * 0.0001 beta_end = scale * 0.02 return np.linspace( beta_start, beta_end, num_diffusion_timesteps, dtype=np.float64 ) elif schedule_name == "cosine": return betas_for_alpha_bar( num_diffusion_timesteps, lambda t: math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2, ) else: raise NotImplementedError(f"unknown beta schedule: {schedule_name}") def betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar, max_beta=0.999): """ Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of (1-beta) over time from t = [0,1]. :param num_diffusion_timesteps: the number of betas to produce. :param alpha_bar: a lambda that takes an argument t from 0 to 1 and produces the cumulative product of (1-beta) up to that part of the diffusion process. :param max_beta: the maximum beta to use; use values lower than 1 to prevent singularities. """ betas = [] for i in range(num_diffusion_timesteps): t1 = i / num_diffusion_timesteps t2 = (i + 1) / num_diffusion_timesteps betas.append(min(1 - alpha_bar(t2) / alpha_bar(t1), max_beta)) return np.array(betas) class ModelMeanType(enum.Enum): """ Which type of output the model predicts. """ PREVIOUS_X = enum.auto() # the model predicts x_{t-1} START_X = enum.auto() # the model predicts x_0 EPSILON = enum.auto() # the model predicts epsilon class ModelVarType(enum.Enum): """ What is used as the model's output variance. The LEARNED_RANGE option has been added to allow the model to predict values between FIXED_SMALL and FIXED_LARGE, making its job easier. """ LEARNED = enum.auto() FIXED_SMALL = enum.auto() FIXED_LARGE = enum.auto() LEARNED_RANGE = enum.auto() class LossType(enum.Enum): MSE = enum.auto() # use raw MSE loss (and KL when learning variances) RESCALED_MSE = ( enum.auto() ) # use raw MSE loss (with RESCALED_KL when learning variances) KL = enum.auto() # use the variational lower-bound RESCALED_KL = enum.auto() # like KL, but rescale to estimate the full VLB def is_vb(self): return self == LossType.KL or self == LossType.RESCALED_KL class GaussianDiffusion: """ Utilities for training and sampling diffusion models. Ported directly from here, and then adapted over time to further experimentation. https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/diffusion_utils_2.py#L42 :param betas: a 1-D numpy array of betas for each diffusion timestep, starting at T and going to 1. :param model_mean_type: a ModelMeanType determining what the model outputs. :param model_var_type: a ModelVarType determining how variance is output. :param loss_type: a LossType determining the loss function to use. :param rescale_timesteps: if True, pass floating point timesteps into the model so that they are always scaled like in the original paper (0 to 1000). """ def __init__( self, *, betas, model_mean_type, model_var_type, loss_type, rescale_timesteps=False, ): self.model_mean_type = model_mean_type self.model_var_type = model_var_type self.loss_type = loss_type self.rescale_timesteps = rescale_timesteps # Use float64 for accuracy. betas = np.array(betas, dtype=np.float64) self.betas = betas assert len(betas.shape) == 1, "betas must be 1-D" assert (betas > 0).all() and (betas <= 1).all() self.num_timesteps = int(betas.shape[0]) alphas = 1.0 - betas self.alphas_cumprod = np.cumprod(alphas, axis=0) self.alphas_cumprod_prev = np.append(1.0, self.alphas_cumprod[:-1]) self.alphas_cumprod_next = np.append(self.alphas_cumprod[1:], 0.0) assert self.alphas_cumprod_prev.shape == (self.num_timesteps,) # calculations for diffusion q(x_t | x_{t-1}) and others self.sqrt_alphas_cumprod = np.sqrt(self.alphas_cumprod) self.sqrt_one_minus_alphas_cumprod = np.sqrt(1.0 - self.alphas_cumprod) self.log_one_minus_alphas_cumprod = np.log(1.0 - self.alphas_cumprod) self.sqrt_recip_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod) self.sqrt_recipm1_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod - 1) # calculations for posterior q(x_{t-1} | x_t, x_0) self.posterior_variance = ( betas * (1.0 - self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod) ) # log calculation clipped because the posterior variance is 0 at the # beginning of the diffusion chain. self.posterior_log_variance_clipped = np.log( np.append(self.posterior_variance[1], self.posterior_variance[1:]) ) self.posterior_mean_coef1 = ( betas * np.sqrt(self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod) ) self.posterior_mean_coef2 = ( (1.0 - self.alphas_cumprod_prev) * np.sqrt(alphas) / (1.0 - self.alphas_cumprod) ) def q_mean_variance(self, x_start, t): """ Get the distribution q(x_t | x_0). :param x_start: the [N x C x ...] tensor of noiseless inputs. :param t: the number of diffusion steps (minus 1). Here, 0 means one step. :return: A tuple (mean, variance, log_variance), all of x_start's shape. """ mean = ( _extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start ) variance = _extract_into_tensor(1.0 - self.alphas_cumprod, t, x_start.shape) log_variance = _extract_into_tensor( self.log_one_minus_alphas_cumprod, t, x_start.shape ) return mean, variance, log_variance def q_sample(self, x_start, t, noise=None): """ Diffuse the data for a given number of diffusion steps. In other words, sample from q(x_t | x_0). :param x_start: the initial data batch. :param t: the number of diffusion steps (minus 1). Here, 0 means one step. :param noise: if specified, the split-out normal noise. :return: A noisy version of x_start. """ raise NotImplementedError if noise is None: noise = th.randn_like(x_start) assert noise.shape == x_start.shape return ( _extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start + _extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise ) def q_sample_blur(self, x_start, t, noise=None): """ Diffuse the data for a given number of diffusion steps. In other words, sample from q(x_t | x_0). :param x_start: the initial data batch. :param t: the number of diffusion steps (minus 1). Here, 0 means one step. :param noise: if specified, the split-out normal noise. :return: A noisy version of x_start. """ assert th.all(t == t[0]) assert float((5*t[0]+1)/self.num_timesteps) >0 , f"t[0]: {t[0]}" out = transforms.functional.gaussian_blur(x_start, 15, float((10*t[0]+1)/self.num_timesteps)) return out def q_posterior_mean_variance(self, x_start, x_t, t): """ Compute the mean and variance of the diffusion posterior: q(x_{t-1} | x_t, x_0) """ assert x_start.shape == x_t.shape posterior_mean = ( _extract_into_tensor(self.posterior_mean_coef1, t, x_t.shape) * x_start + _extract_into_tensor(self.posterior_mean_coef2, t, x_t.shape) * x_t ) posterior_variance = _extract_into_tensor(self.posterior_variance, t, x_t.shape) posterior_log_variance_clipped = _extract_into_tensor( self.posterior_log_variance_clipped, t, x_t.shape ) assert ( posterior_mean.shape[0] == posterior_variance.shape[0] == posterior_log_variance_clipped.shape[0] == x_start.shape[0] ) return posterior_mean, posterior_variance, posterior_log_variance_clipped def p_mean_variance( self, model, x, t, clip_denoised=True, denoised_fn=None, model_kwargs=None ): """ Apply the model to get p(x_{t-1} | x_t), as well as a prediction of the initial x, x_0. :param model: the model, which takes a signal and a batch of timesteps as input. :param x: the [N x C x ...] tensor at time t. :param t: a 1-D Tensor of timesteps. :param clip_denoised: if True, clip the denoised signal into [-1, 1]. :param denoised_fn: if not None, a function which applies to the x_start prediction before it is used to sample. Applies before clip_denoised. :param model_kwargs: if not None, a dict of extra keyword arguments to pass to the model. This can be used for conditioning. :return: a dict with the following keys: - 'mean': the model mean output. - 'variance': the model variance output. - 'log_variance': the log of 'variance'. - 'pred_xstart': the prediction for x_0. """ if model_kwargs is None: model_kwargs = {} B, C = x.shape[:2] assert t.shape == (B,) model_output = model(x, self._scale_timesteps(t), **model_kwargs) if self.model_var_type in [ModelVarType.LEARNED, ModelVarType.LEARNED_RANGE]: assert model_output.shape == (B, C * 2, *x.shape[2:]) model_output, model_var_values = th.split(model_output, C, dim=1) if self.model_var_type == ModelVarType.LEARNED: model_log_variance = model_var_values model_variance = th.exp(model_log_variance) else: min_log = _extract_into_tensor( self.posterior_log_variance_clipped, t, x.shape ) max_log = _extract_into_tensor(np.log(self.betas), t, x.shape) # The model_var_values is [-1, 1] for [min_var, max_var]. frac = (model_var_values + 1) / 2 model_log_variance = frac * max_log + (1 - frac) * min_log model_variance = th.exp(model_log_variance) else: model_variance, model_log_variance = { # for fixedlarge, we set the initial (log-)variance like so # to get a better decoder log likelihood. ModelVarType.FIXED_LARGE: ( np.append(self.posterior_variance[1], self.betas[1:]), np.log(np.append(self.posterior_variance[1], self.betas[1:])), ), ModelVarType.FIXED_SMALL: ( self.posterior_variance, self.posterior_log_variance_clipped, ), }[self.model_var_type] model_variance = _extract_into_tensor(model_variance, t, x.shape) model_log_variance = _extract_into_tensor(model_log_variance, t, x.shape) def process_xstart(x): if denoised_fn is not None: x = denoised_fn(x) if clip_denoised: return x.clamp(-1, 1) return x if self.model_mean_type == ModelMeanType.PREVIOUS_X: pred_xstart = process_xstart( self._predict_xstart_from_xprev(x_t=x, t=t, xprev=model_output) ) model_mean = model_output elif self.model_mean_type in [ModelMeanType.START_X, ModelMeanType.EPSILON]: if self.model_mean_type == ModelMeanType.START_X: pred_xstart = process_xstart(model_output) else: pred_xstart = process_xstart( self._predict_xstart_from_eps(x_t=x, t=t, eps=model_output) ) model_mean, _, _ = self.q_posterior_mean_variance( x_start=pred_xstart, x_t=x, t=t ) else: raise NotImplementedError(self.model_mean_type) assert ( model_mean.shape == model_log_variance.shape == pred_xstart.shape == x.shape ) return { "mean": model_mean, "variance": model_variance, "log_variance": model_log_variance, "pred_xstart": pred_xstart, } def _predict_xstart_from_eps(self, x_t, t, eps): assert x_t.shape == eps.shape return ( _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t - _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * eps ) def _predict_xstart_from_xprev(self, x_t, t, xprev): assert x_t.shape == xprev.shape return ( # (xprev - coef2*x_t) / coef1 _extract_into_tensor(1.0 / self.posterior_mean_coef1, t, x_t.shape) * xprev - _extract_into_tensor( self.posterior_mean_coef2 / self.posterior_mean_coef1, t, x_t.shape ) * x_t ) def _predict_eps_from_xstart(self, x_t, t, pred_xstart): return ( _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t - pred_xstart ) / _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) def _scale_timesteps(self, t): if self.rescale_timesteps: return t.float() * (1000.0 / self.num_timesteps) return t def condition_mean(self, cond_fn, p_mean_var, x, t, model_kwargs=None): """ Compute the mean for the previous step, given a function cond_fn that computes the gradient of a conditional log probability with respect to x. In particular, cond_fn computes grad(log(p(y|x))), and we want to condition on y. This uses the conditioning strategy from Sohl-Dickstein et al. (2015). """ gradient = cond_fn(x, self._scale_timesteps(t), **model_kwargs) new_mean = ( p_mean_var["mean"].float() + p_mean_var["variance"] * gradient.float() ) return new_mean def condition_score(self, cond_fn, p_mean_var, x, t, model_kwargs=None): """ Compute what the p_mean_variance output would have been, should the model's score function be conditioned by cond_fn. See condition_mean() for details on cond_fn. Unlike condition_mean(), this instead uses the conditioning strategy from Song et al (2020). """ alpha_bar = _extract_into_tensor(self.alphas_cumprod, t, x.shape) eps = self._predict_eps_from_xstart(x, t, p_mean_var["pred_xstart"]) eps = eps - (1 - alpha_bar).sqrt() * cond_fn( x, self._scale_timesteps(t), **model_kwargs ) out = p_mean_var.copy() out["pred_xstart"] = self._predict_xstart_from_eps(x, t, eps) out["mean"], _, _ = self.q_posterior_mean_variance( x_start=out["pred_xstart"], x_t=x, t=t ) return out def p_sample( self, model, x, t, clip_denoised=True, denoised_fn=None, cond_fn=None, model_kwargs=None, ): raise NotImplementedError """ Sample x_{t-1} from the model at the given timestep. :param model: the model to sample from. :param x: the current tensor at x_{t-1}. :param t: the value of t, starting at 0 for the first diffusion step. :param clip_denoised: if True, clip the x_start prediction to [-1, 1]. :param denoised_fn: if not None, a function which applies to the x_start prediction before it is used to sample. :param cond_fn: if not None, this is a gradient function that acts similarly to the model. :param model_kwargs: if not None, a dict of extra keyword arguments to pass to the model. This can be used for conditioning. :return: a dict containing the following keys: - 'sample': a random sample from the model. - 'pred_xstart': a prediction of x_0. """ out = self.p_mean_variance( model, x, t, clip_denoised=clip_denoised, denoised_fn=denoised_fn, model_kwargs=model_kwargs, ) noise = th.randn_like(x) nonzero_mask = ( (t != 0).float().view(-1, *([1] * (len(x.shape) - 1))) ) # no noise when t == 0 if cond_fn is not None: out["mean"] = self.condition_mean( cond_fn, out, x, t, model_kwargs=model_kwargs ) sample = out["mean"] + nonzero_mask * th.exp(0.5 * out["log_variance"]) * noise return {"sample": sample, "pred_xstart": out["pred_xstart"]} def p_sample_blur( self, model, x, t, clip_denoised=True, denoised_fn=None, cond_fn=None, model_kwargs=None, ): """ Sample x_{t-1} from the model at the given timestep. :param model: the model to sample from. :param x: the current tensor at x_{t-1}. :param t: the value of t, starting at 0 for the first diffusion step. :param clip_denoised: if True, clip the x_start prediction to [-1, 1]. :param denoised_fn: if not None, a function which applies to the x_start prediction before it is used to sample. :param cond_fn: if not None, this is a gradient function that acts similarly to the model. :param model_kwargs: if not None, a dict of extra keyword arguments to pass to the model. This can be used for conditioning. :return: a dict containing the following keys: - 'sample': a random sample from the model. - 'pred_xstart': a prediction of x_0. """ # out = self.p_mean_variance( # model, # x, # t, # clip_denoised=clip_denoised, # denoised_fn=denoised_fn, # model_kwargs=model_kwargs, # ) # noise = th.randn_like(x) # nonzero_mask = ( # (t != 0).float().view(-1, *([1] * (len(x.shape) - 1))) # ) # no noise when t == 0 if cond_fn is not None: raise NotImplementedError out["mean"] = self.condition_mean( cond_fn, out, x, t, model_kwargs=model_kwargs ) model_output = model(x, self._scale_timesteps(t), **model_kwargs) # sample = out["mean"] + nonzero_mask * th.exp(0.5 * out["log_variance"]) * noise return {"sample": model_output} def p_sample_loop( self, model, shape, noise=None, clip_denoised=True, denoised_fn=None, cond_fn=None, model_kwargs=None, device=None, progress=False, ): """ Generate samples from the model. :param model: the model module. :param shape: the shape of the samples, (N, C, H, W). :param noise: if specified, the noise from the encoder to sample. Should be of the same shape as `shape`. :param clip_denoised: if True, clip x_start predictions to [-1, 1]. :param denoised_fn: if not None, a function which applies to the x_start prediction before it is used to sample. :param cond_fn: if not None, this is a gradient function that acts similarly to the model. :param model_kwargs: if not None, a dict of extra keyword arguments to pass to the model. This can be used for conditioning. :param device: if specified, the device to create the samples on. If not specified, use a model parameter's device. :param progress: if True, show a tqdm progress bar. :return: a non-differentiable batch of samples. """ final = None i=0 for sample in self.p_sample_loop_progressive( model, shape, noise=noise, clip_denoised=clip_denoised, denoised_fn=denoised_fn, cond_fn=cond_fn, model_kwargs=model_kwargs, device=device, progress=progress, ): print("p_sample_loop", i) i+=1 tensor_imsave(sample["sample"][0], "./debug", f"{self.num_timesteps-i}.png", prt=True) final = sample return final["sample"] def p_sample_loop_progressive( self, model, shape, noise=None, clip_denoised=True, denoised_fn=None, cond_fn=None, model_kwargs=None, device=None, progress=False, ): """ Generate samples from the model and yield intermediate samples from each timestep of diffusion. Arguments are the same as p_sample_loop(). Returns a generator over dicts, where each dict is the return value of p_sample(). """ if device is None: device = next(model.parameters()).device assert isinstance(shape, (tuple, list)) if noise is not None: img = noise else: img = th.randn(*shape, device=device) T = th.tensor(1000, device=device).repeat(shape[0]) print("T:", T) img = self.q_sample_blur(img, T) indices = list(range(self.num_timesteps))[::-1] if progress: # Lazy import so that we don't depend on tqdm. from tqdm.auto import tqdm indices = tqdm(indices) for i in indices: t = th.tensor([i] * shape[0], device=device) with th.no_grad(): out = self.p_sample_blur( model, img, t, clip_denoised=clip_denoised, denoised_fn=denoised_fn, cond_fn=cond_fn, model_kwargs=model_kwargs, ) yield out img = out["sample"] def ddim_sample( self, model, x, t, clip_denoised=True, denoised_fn=None, cond_fn=None, model_kwargs=None, eta=0.0, ): """ Sample x_{t-1} from the model using DDIM. Same usage as p_sample(). """ out = self.p_mean_variance( model, x, t, clip_denoised=clip_denoised, denoised_fn=denoised_fn, model_kwargs=model_kwargs, ) if cond_fn is not None: out = self.condition_score(cond_fn, out, x, t, model_kwargs=model_kwargs) # Usually our model outputs epsilon, but we re-derive it # in case we used x_start or x_prev prediction. eps = self._predict_eps_from_xstart(x, t, out["pred_xstart"]) alpha_bar = _extract_into_tensor(self.alphas_cumprod, t, x.shape) alpha_bar_prev = _extract_into_tensor(self.alphas_cumprod_prev, t, x.shape) sigma = ( eta * th.sqrt((1 - alpha_bar_prev) / (1 - alpha_bar)) * th.sqrt(1 - alpha_bar / alpha_bar_prev) ) # Equation 12. noise = th.randn_like(x) mean_pred = ( out["pred_xstart"] * th.sqrt(alpha_bar_prev) + th.sqrt(1 - alpha_bar_prev - sigma ** 2) * eps ) nonzero_mask = ( (t != 0).float().view(-1, *([1] * (len(x.shape) - 1))) ) # no noise when t == 0 sample = mean_pred + nonzero_mask * sigma * noise return {"sample": sample, "pred_xstart": out["pred_xstart"]} def ddim_reverse_sample( self, model, x, t, clip_denoised=True, denoised_fn=None, model_kwargs=None, eta=0.0, ): """ Sample x_{t+1} from the model using DDIM reverse ODE. """ assert eta == 0.0, "Reverse ODE only for deterministic path" out = self.p_mean_variance( model, x, t, clip_denoised=clip_denoised, denoised_fn=denoised_fn, model_kwargs=model_kwargs, ) # Usually our model outputs epsilon, but we re-derive it # in case we used x_start or x_prev prediction. eps = ( _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x.shape) * x - out["pred_xstart"] ) / _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x.shape) alpha_bar_next = _extract_into_tensor(self.alphas_cumprod_next, t, x.shape) # Equation 12. reversed mean_pred = ( out["pred_xstart"] * th.sqrt(alpha_bar_next) + th.sqrt(1 - alpha_bar_next) * eps ) return {"sample": mean_pred, "pred_xstart": out["pred_xstart"]} def ddim_sample_loop( self, model, shape, noise=None, clip_denoised=True, denoised_fn=None, cond_fn=None, model_kwargs=None, device=None, progress=False, eta=0.0, ): """ Generate samples from the model using DDIM. Same usage as p_sample_loop(). """ final = None for sample in self.ddim_sample_loop_progressive( model, shape, noise=noise, clip_denoised=clip_denoised, denoised_fn=denoised_fn, cond_fn=cond_fn, model_kwargs=model_kwargs, device=device, progress=progress, eta=eta, ): final = sample return final["sample"] def ddim_sample_loop_progressive( self, model, shape, noise=None, clip_denoised=True, denoised_fn=None, cond_fn=None, model_kwargs=None, device=None, progress=False, eta=0.0, ): """ Use DDIM to sample from the model and yield intermediate samples from each timestep of DDIM. Same usage as p_sample_loop_progressive(). """ if device is None: device = next(model.parameters()).device assert isinstance(shape, (tuple, list)) if noise is not None: img = noise else: img = th.randn(*shape, device=device) indices = list(range(self.num_timesteps))[::-1] if progress: # Lazy import so that we don't depend on tqdm. from tqdm.auto import tqdm indices = tqdm(indices) for i in indices: t = th.tensor([i] * shape[0], device=device) with th.no_grad(): out = self.ddim_sample( model, img, t, clip_denoised=clip_denoised, denoised_fn=denoised_fn, cond_fn=cond_fn, model_kwargs=model_kwargs, eta=eta, ) yield out img = out["sample"] def _vb_terms_bpd( self, model, x_start, x_t, t, clip_denoised=True, model_kwargs=None ): """ Get a term for the variational lower-bound. The resulting units are bits (rather than nats, as one might expect). This allows for comparison to other papers. :return: a dict with the following keys: - 'output': a shape [N] tensor of NLLs or KLs. - 'pred_xstart': the x_0 predictions. """ true_mean, _, true_log_variance_clipped = self.q_posterior_mean_variance( x_start=x_start, x_t=x_t, t=t ) out = self.p_mean_variance( model, x_t, t, clip_denoised=clip_denoised, model_kwargs=model_kwargs ) kl = normal_kl( true_mean, true_log_variance_clipped, out["mean"], out["log_variance"] ) kl = mean_flat(kl) / np.log(2.0) decoder_nll = -discretized_gaussian_log_likelihood( x_start, means=out["mean"], log_scales=0.5 * out["log_variance"] ) assert decoder_nll.shape == x_start.shape decoder_nll = mean_flat(decoder_nll) / np.log(2.0) # At the first timestep return the decoder NLL, # otherwise return KL(q(x_{t-1}|x_t,x_0) || p(x_{t-1}|x_t)) output = th.where((t == 0), decoder_nll, kl) return {"output": output, "pred_xstart": out["pred_xstart"]} def training_losses(self, model, x_start, t, model_kwargs=None, noise=None): """ Compute training losses for a single timestep. :param model: the model to evaluate loss on. :param x_start: the [N x C x ...] tensor of inputs. :param t: a batch of timestep indices. :param model_kwargs: if not None, a dict of extra keyword arguments to pass to the model. This can be used for conditioning. :param noise: if specified, the specific Gaussian noise to try to remove. :return: a dict with the key "loss" containing a tensor of shape [N]. Some mean or variance settings may also have other keys. """ if model_kwargs is None: model_kwargs = {} if noise is None: noise = th.randn_like(x_start) # x_t = self.q_sample(x_start, t, noise=noise) x_t = self.q_sample_blur(x_start, t) x_t_prev = self.q_sample_blur(x_start, t-1) # x_{t-1} terms = {} if self.loss_type == LossType.KL or self.loss_type == LossType.RESCALED_KL: raise NotImplementedError terms["loss"] = self._vb_terms_bpd( model=model, x_start=x_start, x_t=x_t, t=t, clip_denoised=False, model_kwargs=model_kwargs, )["output"] if self.loss_type == LossType.RESCALED_KL: terms["loss"] *= self.num_timesteps elif self.loss_type == LossType.MSE or self.loss_type == LossType.RESCALED_MSE: model_output = model(x_t, self._scale_timesteps(t), **model_kwargs) if self.model_var_type in [ ModelVarType.LEARNED, ModelVarType.LEARNED_RANGE, ]: raise NotImplementedError B, C = x_t.shape[:2] assert model_output.shape == (B, C * 2, *x_t.shape[2:]) model_output, model_var_values = th.split(model_output, C, dim=1) # Learn the variance using the variational bound, but don't let # it affect our mean prediction. frozen_out = th.cat([model_output.detach(), model_var_values], dim=1) terms["vb"] = self._vb_terms_bpd( model=lambda *args, r=frozen_out: r, x_start=x_start, x_t=x_t, t=t, clip_denoised=False, )["output"] if self.loss_type == LossType.RESCALED_MSE: # Divide by 1000 for equivalence with initial implementation. # Without a factor of 1/1000, the VB term hurts the MSE term. terms["vb"] *= self.num_timesteps / 1000.0 # target = { # ModelMeanType.PREVIOUS_X: self.q_posterior_mean_variance( # x_start=x_start, x_t=x_t, t=t # )[0], # ModelMeanType.START_X: x_start, # ModelMeanType.EPSILON: noise, # }[self.model_mean_type] target = x_t_prev assert model_output.shape == target.shape == x_start.shape terms["mse"] = mean_flat((target - model_output) ** 2) if "vb" in terms: terms["loss"] = terms["mse"] + terms["vb"] else: terms["loss"] = terms["mse"] else: raise NotImplementedError(self.loss_type) return terms def _prior_bpd(self, x_start): """ Get the prior KL term for the variational lower-bound, measured in bits-per-dim. This term can't be optimized, as it only depends on the encoder. :param x_start: the [N x C x ...] tensor of inputs. :return: a batch of [N] KL values (in bits), one per batch element. """ batch_size = x_start.shape[0] t = th.tensor([self.num_timesteps - 1] * batch_size, device=x_start.device) qt_mean, _, qt_log_variance = self.q_mean_variance(x_start, t) kl_prior = normal_kl( mean1=qt_mean, logvar1=qt_log_variance, mean2=0.0, logvar2=0.0 ) return mean_flat(kl_prior) / np.log(2.0) def calc_bpd_loop(self, model, x_start, clip_denoised=True, model_kwargs=None): """ Compute the entire variational lower-bound, measured in bits-per-dim, as well as other related quantities. :param model: the model to evaluate loss on. :param x_start: the [N x C x ...] tensor of inputs. :param clip_denoised: if True, clip denoised samples. :param model_kwargs: if not None, a dict of extra keyword arguments to pass to the model. This can be used for conditioning. :return: a dict containing the following keys: - total_bpd: the total variational lower-bound, per batch element. - prior_bpd: the prior term in the lower-bound. - vb: an [N x T] tensor of terms in the lower-bound. - xstart_mse: an [N x T] tensor of x_0 MSEs for each timestep. - mse: an [N x T] tensor of epsilon MSEs for each timestep. """ device = x_start.device batch_size = x_start.shape[0] xstart_mse = [] mse = [] for t in list(range(self.num_timesteps))[::-1]: t_batch = th.tensor([t] * batch_size, device=device) x_t = self.q_sample_blur(x_start=x_start, t=t_batch) x_t_prev = self.q_sample_blur(x_start=x_start, t=t_batch-1) # Calculate VLB term at the current timestep with th.no_grad(): model_output = model(x_t, self._scale_timesteps(t), **model_kwargs) mse.append(mean_flat((x_t_prev - model_output)**2)) mse = th.stack(mse, dim=1) prior_bpd = self._prior_bpd(x_start) return {'mse' : mse} # device = x_start.device # batch_size = x_start.shape[0] # vb = [] # xstart_mse = [] # mse = [] # for t in list(range(self.num_timesteps))[::-1]: # t_batch = th.tensor([t] * batch_size, device=device) # noise = th.randn_like(x_start) # x_t = self.q_sample(x_start=x_start, t=t_batch, noise=noise) # # Calculate VLB term at the current timestep # with th.no_grad(): # out = self._vb_terms_bpd( # model, # x_start=x_start, # x_t=x_t, # t=t_batch, # clip_denoised=clip_denoised, # model_kwargs=model_kwargs, # ) # vb.append(out["output"]) # xstart_mse.append(mean_flat((out["pred_xstart"] - x_start) ** 2)) # eps = self._predict_eps_from_xstart(x_t, t_batch, out["pred_xstart"]) # mse.append(mean_flat((eps - noise) ** 2)) # vb = th.stack(vb, dim=1) # xstart_mse = th.stack(xstart_mse, dim=1) # mse = th.stack(mse, dim=1) # prior_bpd = self._prior_bpd(x_start) # total_bpd = vb.sum(dim=1) + prior_bpd # return { # "total_bpd": total_bpd, # "prior_bpd": prior_bpd, # "vb": vb, # "xstart_mse": xstart_mse, # "mse": mse, # } def _extract_into_tensor(arr, timesteps, broadcast_shape): """ Extract values from a 1-D numpy array for a batch of indices. :param arr: the 1-D numpy array. :param timesteps: a tensor of indices into the array to extract. :param broadcast_shape: a larger shape of K dimensions with the batch dimension equal to the length of timesteps. :return: a tensor of shape [batch_size, 1, ...] where the shape has K dims. """ res = th.from_numpy(arr).to(device=timesteps.device)[timesteps].float() while len(res.shape) < len(broadcast_shape): res = res[..., None] return res.expand(broadcast_shape) def tensor2pil(t): img = transforms.functional.to_pil_image(denorm(t)) return img def denorm(t): return t*0.5 + 0.5 def tensor_imsave(t, path, fname, denormalization=True, prt=False): # Save tensor as .png file check_folder(path) if denormalization: t = denorm(t) t = th.clamp(input=t, min=0, max=1) img = transforms.functional.to_pil_image(t.detach().cpu()) img.save(os.path.join(path,fname)) if prt: print(f"Saved to {os.path.join(path,fname)}") def pil_loader(path): # open path as file to avoid ResourceWarning (https://github.com/python-pillow/Pillow/issues/835) with open(path, 'rb') as f:##Automatic closing using with. rb = read binary img = Image.open(f) return img.convert('RGB') def check_folder(log_dir): if not os.path.exists(log_dir): os.makedirs(log_dir) return log_dir ================================================ FILE: guided_diffusion/image_datasets.py ================================================ import math import random from PIL import Image import blobfile as bf from mpi4py import MPI import numpy as np from torch.utils.data import DataLoader, Dataset def load_data( *, data_dir, batch_size, image_size, class_cond=False, deterministic=False, random_crop=False, random_flip=True, ): """ For a dataset, create a generator over (images, kwargs) pairs. Each images is an NCHW float tensor, and the kwargs dict contains zero or more keys, each of which map to a batched Tensor of their own. The kwargs dict can be used for class labels, in which case the key is "y" and the values are integer tensors of class labels. :param data_dir: a dataset directory. :param batch_size: the batch size of each returned pair. :param image_size: the size to which images are resized. :param class_cond: if True, include a "y" key in returned dicts for class label. If classes are not available and this is true, an exception will be raised. :param deterministic: if True, yield results in a deterministic order. :param random_crop: if True, randomly crop the images for augmentation. :param random_flip: if True, randomly flip the images for augmentation. """ if not data_dir: raise ValueError("unspecified data directory") all_files = _list_image_files_recursively(data_dir) classes = None if class_cond: # Assume classes are the first part of the filename, # before an underscore. class_names = [bf.basename(path).split("_")[0] for path in all_files] sorted_classes = {x: i for i, x in enumerate(sorted(set(class_names)))} classes = [sorted_classes[x] for x in class_names] dataset = ImageDataset( image_size, all_files, classes=classes, shard=MPI.COMM_WORLD.Get_rank(), num_shards=MPI.COMM_WORLD.Get_size(), random_crop=random_crop, random_flip=random_flip, ) if deterministic: loader = DataLoader( dataset, batch_size=batch_size, shuffle=False, num_workers=1, drop_last=True ) else: loader = DataLoader( dataset, batch_size=batch_size, shuffle=True, num_workers=1, drop_last=True ) while True: yield from loader def _list_image_files_recursively(data_dir): results = [] for entry in sorted(bf.listdir(data_dir)): full_path = bf.join(data_dir, entry) ext = entry.split(".")[-1] if "." in entry and ext.lower() in ["jpg", "jpeg", "png", "gif"]: results.append(full_path) elif bf.isdir(full_path): results.extend(_list_image_files_recursively(full_path)) return results class ImageDataset(Dataset): def __init__( self, resolution, image_paths, classes=None, shard=0, num_shards=1, random_crop=False, random_flip=True, ): super().__init__() self.resolution = resolution self.local_images = image_paths[shard:][::num_shards] self.local_classes = None if classes is None else classes[shard:][::num_shards] self.random_crop = random_crop self.random_flip = random_flip def __len__(self): return len(self.local_images) def __getitem__(self, idx): path = self.local_images[idx] with bf.BlobFile(path, "rb") as f: pil_image = Image.open(f) pil_image.load() pil_image = pil_image.convert("RGB") if self.random_crop: arr = random_crop_arr(pil_image, self.resolution) else: arr = center_crop_arr(pil_image, self.resolution) if self.random_flip and random.random() < 0.5: arr = arr[:, ::-1] arr = arr.astype(np.float32) / 127.5 - 1 out_dict = {} if self.local_classes is not None: out_dict["y"] = np.array(self.local_classes[idx], dtype=np.int64) return np.transpose(arr, [2, 0, 1]), out_dict def center_crop_arr(pil_image, image_size): # We are not on a new enough PIL to support the `reducing_gap` # argument, which uses BOX downsampling at powers of two first. # Thus, we do it by hand to improve downsample quality. while min(*pil_image.size) >= 2 * image_size: pil_image = pil_image.resize( tuple(x // 2 for x in pil_image.size), resample=Image.BOX ) scale = image_size / min(*pil_image.size) pil_image = pil_image.resize( tuple(round(x * scale) for x in pil_image.size), resample=Image.BICUBIC ) arr = np.array(pil_image) crop_y = (arr.shape[0] - image_size) // 2 crop_x = (arr.shape[1] - image_size) // 2 return arr[crop_y : crop_y + image_size, crop_x : crop_x + image_size] def random_crop_arr(pil_image, image_size, min_crop_frac=0.8, max_crop_frac=1.0): min_smaller_dim_size = math.ceil(image_size / max_crop_frac) max_smaller_dim_size = math.ceil(image_size / min_crop_frac) smaller_dim_size = random.randrange(min_smaller_dim_size, max_smaller_dim_size + 1) # We are not on a new enough PIL to support the `reducing_gap` # argument, which uses BOX downsampling at powers of two first. # Thus, we do it by hand to improve downsample quality. while min(*pil_image.size) >= 2 * smaller_dim_size: pil_image = pil_image.resize( tuple(x // 2 for x in pil_image.size), resample=Image.BOX ) scale = smaller_dim_size / min(*pil_image.size) pil_image = pil_image.resize( tuple(round(x * scale) for x in pil_image.size), resample=Image.BICUBIC ) arr = np.array(pil_image) crop_y = random.randrange(arr.shape[0] - image_size + 1) crop_x = random.randrange(arr.shape[1] - image_size + 1) return arr[crop_y : crop_y + image_size, crop_x : crop_x + image_size] ================================================ FILE: guided_diffusion/logger.py ================================================ """ Logger copied from OpenAI baselines to avoid extra RL-based dependencies: https://github.com/openai/baselines/blob/ea25b9e8b234e6ee1bca43083f8f3cf974143998/baselines/logger.py """ import os import sys import shutil import os.path as osp import json import time import datetime import tempfile import warnings from collections import defaultdict from contextlib import contextmanager DEBUG = 10 INFO = 20 WARN = 30 ERROR = 40 DISABLED = 50 class KVWriter(object): def writekvs(self, kvs): raise NotImplementedError class SeqWriter(object): def writeseq(self, seq): raise NotImplementedError class HumanOutputFormat(KVWriter, SeqWriter): def __init__(self, filename_or_file): if isinstance(filename_or_file, str): self.file = open(filename_or_file, "wt") self.own_file = True else: assert hasattr(filename_or_file, "read"), ( "expected file or str, got %s" % filename_or_file ) self.file = filename_or_file self.own_file = False def writekvs(self, kvs): # Create strings for printing key2str = {} for (key, val) in sorted(kvs.items()): if hasattr(val, "__float__"): valstr = "%-8.3g" % val else: valstr = str(val) key2str[self._truncate(key)] = self._truncate(valstr) # Find max widths if len(key2str) == 0: print("WARNING: tried to write empty key-value dict") return else: keywidth = max(map(len, key2str.keys())) valwidth = max(map(len, key2str.values())) # Write out the data dashes = "-" * (keywidth + valwidth + 7) lines = [dashes] for (key, val) in sorted(key2str.items(), key=lambda kv: kv[0].lower()): lines.append( "| %s%s | %s%s |" % (key, " " * (keywidth - len(key)), val, " " * (valwidth - len(val))) ) lines.append(dashes) self.file.write("\n".join(lines) + "\n") # Flush the output to the file self.file.flush() def _truncate(self, s): maxlen = 30 return s[: maxlen - 3] + "..." if len(s) > maxlen else s def writeseq(self, seq): seq = list(seq) for (i, elem) in enumerate(seq): self.file.write(elem) if i < len(seq) - 1: # add space unless this is the last one self.file.write(" ") self.file.write("\n") self.file.flush() def close(self): if self.own_file: self.file.close() class JSONOutputFormat(KVWriter): def __init__(self, filename): self.file = open(filename, "wt") def writekvs(self, kvs): for k, v in sorted(kvs.items()): if hasattr(v, "dtype"): kvs[k] = float(v) self.file.write(json.dumps(kvs) + "\n") self.file.flush() def close(self): self.file.close() class CSVOutputFormat(KVWriter): def __init__(self, filename): self.file = open(filename, "w+t") self.keys = [] self.sep = "," def writekvs(self, kvs): # Add our current row to the history extra_keys = list(kvs.keys() - self.keys) extra_keys.sort() if extra_keys: self.keys.extend(extra_keys) self.file.seek(0) lines = self.file.readlines() self.file.seek(0) for (i, k) in enumerate(self.keys): if i > 0: self.file.write(",") self.file.write(k) self.file.write("\n") for line in lines[1:]: self.file.write(line[:-1]) self.file.write(self.sep * len(extra_keys)) self.file.write("\n") for (i, k) in enumerate(self.keys): if i > 0: self.file.write(",") v = kvs.get(k) if v is not None: self.file.write(str(v)) self.file.write("\n") self.file.flush() def close(self): self.file.close() class TensorBoardOutputFormat(KVWriter): """ Dumps key/value pairs into TensorBoard's numeric format. """ def __init__(self, dir): os.makedirs(dir, exist_ok=True) self.dir = dir self.step = 1 prefix = "events" path = osp.join(osp.abspath(dir), prefix) import tensorflow as tf from tensorflow.python import pywrap_tensorflow from tensorflow.core.util import event_pb2 from tensorflow.python.util import compat self.tf = tf self.event_pb2 = event_pb2 self.pywrap_tensorflow = pywrap_tensorflow self.writer = pywrap_tensorflow.EventsWriter(compat.as_bytes(path)) def writekvs(self, kvs): def summary_val(k, v): kwargs = {"tag": k, "simple_value": float(v)} return self.tf.Summary.Value(**kwargs) summary = self.tf.Summary(value=[summary_val(k, v) for k, v in kvs.items()]) event = self.event_pb2.Event(wall_time=time.time(), summary=summary) event.step = ( self.step ) # is there any reason why you'd want to specify the step? self.writer.WriteEvent(event) self.writer.Flush() self.step += 1 def close(self): if self.writer: self.writer.Close() self.writer = None def make_output_format(format, ev_dir, log_suffix=""): os.makedirs(ev_dir, exist_ok=True) if format == "stdout": return HumanOutputFormat(sys.stdout) elif format == "log": return HumanOutputFormat(osp.join(ev_dir, "log%s.txt" % log_suffix)) elif format == "json": return JSONOutputFormat(osp.join(ev_dir, "progress%s.json" % log_suffix)) elif format == "csv": return CSVOutputFormat(osp.join(ev_dir, "progress%s.csv" % log_suffix)) elif format == "tensorboard": return TensorBoardOutputFormat(osp.join(ev_dir, "tb%s" % log_suffix)) else: raise ValueError("Unknown format specified: %s" % (format,)) # ================================================================ # API # ================================================================ def logkv(key, val): """ Log a value of some diagnostic Call this once for each diagnostic quantity, each iteration If called many times, last value will be used. """ get_current().logkv(key, val) def logkv_mean(key, val): """ The same as logkv(), but if called many times, values averaged. """ get_current().logkv_mean(key, val) def logkvs(d): """ Log a dictionary of key-value pairs """ for (k, v) in d.items(): logkv(k, v) def dumpkvs(): """ Write all of the diagnostics from the current iteration """ return get_current().dumpkvs() def getkvs(): return get_current().name2val def log(*args, level=INFO): """ Write the sequence of args, with no separators, to the console and output files (if you've configured an output file). """ get_current().log(*args, level=level) def debug(*args): log(*args, level=DEBUG) def info(*args): log(*args, level=INFO) def warn(*args): log(*args, level=WARN) def error(*args): log(*args, level=ERROR) def set_level(level): """ Set logging threshold on current logger. """ get_current().set_level(level) def set_comm(comm): get_current().set_comm(comm) def get_dir(): """ Get directory that log files are being written to. will be None if there is no output directory (i.e., if you didn't call start) """ return get_current().get_dir() record_tabular = logkv dump_tabular = dumpkvs @contextmanager def profile_kv(scopename): logkey = "wait_" + scopename tstart = time.time() try: yield finally: get_current().name2val[logkey] += time.time() - tstart def profile(n): """ Usage: @profile("my_func") def my_func(): code """ def decorator_with_name(func): def func_wrapper(*args, **kwargs): with profile_kv(n): return func(*args, **kwargs) return func_wrapper return decorator_with_name # ================================================================ # Backend # ================================================================ def get_current(): if Logger.CURRENT is None: _configure_default_logger() return Logger.CURRENT class Logger(object): DEFAULT = None # A logger with no output files. (See right below class definition) # So that you can still log to the terminal without setting up any output files CURRENT = None # Current logger being used by the free functions above def __init__(self, dir, output_formats, comm=None): self.name2val = defaultdict(float) # values this iteration self.name2cnt = defaultdict(int) self.level = INFO self.dir = dir self.output_formats = output_formats self.comm = comm # Logging API, forwarded # ---------------------------------------- def logkv(self, key, val): self.name2val[key] = val def logkv_mean(self, key, val): oldval, cnt = self.name2val[key], self.name2cnt[key] self.name2val[key] = oldval * cnt / (cnt + 1) + val / (cnt + 1) self.name2cnt[key] = cnt + 1 def dumpkvs(self): if self.comm is None: d = self.name2val else: d = mpi_weighted_mean( self.comm, { name: (val, self.name2cnt.get(name, 1)) for (name, val) in self.name2val.items() }, ) if self.comm.rank != 0: d["dummy"] = 1 # so we don't get a warning about empty dict out = d.copy() # Return the dict for unit testing purposes for fmt in self.output_formats: if isinstance(fmt, KVWriter): fmt.writekvs(d) self.name2val.clear() self.name2cnt.clear() return out def log(self, *args, level=INFO): if self.level <= level: self._do_log(args) # Configuration # ---------------------------------------- def set_level(self, level): self.level = level def set_comm(self, comm): self.comm = comm def get_dir(self): return self.dir def close(self): for fmt in self.output_formats: fmt.close() # Misc # ---------------------------------------- def _do_log(self, args): for fmt in self.output_formats: if isinstance(fmt, SeqWriter): fmt.writeseq(map(str, args)) def get_rank_without_mpi_import(): # check environment variables here instead of importing mpi4py # to avoid calling MPI_Init() when this module is imported for varname in ["PMI_RANK", "OMPI_COMM_WORLD_RANK"]: if varname in os.environ: return int(os.environ[varname]) return 0 def mpi_weighted_mean(comm, local_name2valcount): """ Copied from: https://github.com/openai/baselines/blob/ea25b9e8b234e6ee1bca43083f8f3cf974143998/baselines/common/mpi_util.py#L110 Perform a weighted average over dicts that are each on a different node Input: local_name2valcount: dict mapping key -> (value, count) Returns: key -> mean """ all_name2valcount = comm.gather(local_name2valcount) if comm.rank == 0: name2sum = defaultdict(float) name2count = defaultdict(float) for n2vc in all_name2valcount: for (name, (val, count)) in n2vc.items(): try: val = float(val) except ValueError: if comm.rank == 0: warnings.warn( "WARNING: tried to compute mean on non-float {}={}".format( name, val ) ) else: name2sum[name] += val * count name2count[name] += count return {name: name2sum[name] / name2count[name] for name in name2sum} else: return {} def configure(dir=None, format_strs=None, comm=None, log_suffix=""): """ If comm is provided, average all numerical stats across that comm """ dir = "./cifar10-forwardblur-fromblurrednoise" if dir is None: dir = os.getenv("OPENAI_LOGDIR") if dir is None: dir = osp.join( tempfile.gettempdir(), datetime.datetime.now().strftime("openai-%Y-%m-%d-%H-%M-%S-%f"), ) assert isinstance(dir, str) dir = os.path.expanduser(dir) os.makedirs(os.path.expanduser(dir), exist_ok=True) rank = get_rank_without_mpi_import() if rank > 0: log_suffix = log_suffix + "-rank%03i" % rank if format_strs is None: if rank == 0: format_strs = os.getenv("OPENAI_LOG_FORMAT", "stdout,log,csv").split(",") else: format_strs = os.getenv("OPENAI_LOG_FORMAT_MPI", "log").split(",") format_strs = filter(None, format_strs) output_formats = [make_output_format(f, dir, log_suffix) for f in format_strs] Logger.CURRENT = Logger(dir=dir, output_formats=output_formats, comm=comm) if output_formats: log("Logging to %s" % dir) def _configure_default_logger(): configure() Logger.DEFAULT = Logger.CURRENT def reset(): if Logger.CURRENT is not Logger.DEFAULT: Logger.CURRENT.close() Logger.CURRENT = Logger.DEFAULT log("Reset logger") @contextmanager def scoped_configure(dir=None, format_strs=None, comm=None): prevlogger = Logger.CURRENT configure(dir=dir, format_strs=format_strs, comm=comm) try: yield finally: Logger.CURRENT.close() Logger.CURRENT = prevlogger ================================================ FILE: guided_diffusion/losses.py ================================================ """ Helpers for various likelihood-based losses. These are ported from the original Ho et al. diffusion models codebase: https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/utils.py """ import numpy as np import torch as th def normal_kl(mean1, logvar1, mean2, logvar2): """ Compute the KL divergence between two gaussians. Shapes are automatically broadcasted, so batches can be compared to scalars, among other use cases. """ tensor = None for obj in (mean1, logvar1, mean2, logvar2): if isinstance(obj, th.Tensor): tensor = obj break assert tensor is not None, "at least one argument must be a Tensor" # Force variances to be Tensors. Broadcasting helps convert scalars to # Tensors, but it does not work for th.exp(). logvar1, logvar2 = [ x if isinstance(x, th.Tensor) else th.tensor(x).to(tensor) for x in (logvar1, logvar2) ] return 0.5 * ( -1.0 + logvar2 - logvar1 + th.exp(logvar1 - logvar2) + ((mean1 - mean2) ** 2) * th.exp(-logvar2) ) def approx_standard_normal_cdf(x): """ A fast approximation of the cumulative distribution function of the standard normal. """ return 0.5 * (1.0 + th.tanh(np.sqrt(2.0 / np.pi) * (x + 0.044715 * th.pow(x, 3)))) def discretized_gaussian_log_likelihood(x, *, means, log_scales): """ Compute the log-likelihood of a Gaussian distribution discretizing to a given image. :param x: the target images. It is assumed that this was uint8 values, rescaled to the range [-1, 1]. :param means: the Gaussian mean Tensor. :param log_scales: the Gaussian log stddev Tensor. :return: a tensor like x of log probabilities (in nats). """ assert x.shape == means.shape == log_scales.shape centered_x = x - means inv_stdv = th.exp(-log_scales) plus_in = inv_stdv * (centered_x + 1.0 / 255.0) cdf_plus = approx_standard_normal_cdf(plus_in) min_in = inv_stdv * (centered_x - 1.0 / 255.0) cdf_min = approx_standard_normal_cdf(min_in) log_cdf_plus = th.log(cdf_plus.clamp(min=1e-12)) log_one_minus_cdf_min = th.log((1.0 - cdf_min).clamp(min=1e-12)) cdf_delta = cdf_plus - cdf_min log_probs = th.where( x < -0.999, log_cdf_plus, th.where(x > 0.999, log_one_minus_cdf_min, th.log(cdf_delta.clamp(min=1e-12))), ) assert log_probs.shape == x.shape return log_probs ================================================ FILE: guided_diffusion/nn.py ================================================ """ Various utilities for neural networks. """ import math import torch as th import torch.nn as nn # PyTorch 1.7 has SiLU, but we support PyTorch 1.5. class SiLU(nn.Module): def forward(self, x): return x * th.sigmoid(x) class GroupNorm32(nn.GroupNorm): def forward(self, x): return super().forward(x.float()).type(x.dtype) def conv_nd(dims, *args, **kwargs): """ Create a 1D, 2D, or 3D convolution module. """ if dims == 1: return nn.Conv1d(*args, **kwargs) elif dims == 2: return nn.Conv2d(*args, **kwargs) elif dims == 3: return nn.Conv3d(*args, **kwargs) raise ValueError(f"unsupported dimensions: {dims}") def linear(*args, **kwargs): """ Create a linear module. """ return nn.Linear(*args, **kwargs) def avg_pool_nd(dims, *args, **kwargs): """ Create a 1D, 2D, or 3D average pooling module. """ if dims == 1: return nn.AvgPool1d(*args, **kwargs) elif dims == 2: return nn.AvgPool2d(*args, **kwargs) elif dims == 3: return nn.AvgPool3d(*args, **kwargs) raise ValueError(f"unsupported dimensions: {dims}") def update_ema(target_params, source_params, rate=0.99): """ Update target parameters to be closer to those of source parameters using an exponential moving average. :param target_params: the target parameter sequence. :param source_params: the source parameter sequence. :param rate: the EMA rate (closer to 1 means slower). """ for targ, src in zip(target_params, source_params): targ.detach().mul_(rate).add_(src, alpha=1 - rate) def zero_module(module): """ Zero out the parameters of a module and return it. """ for p in module.parameters(): p.detach().zero_() return module def scale_module(module, scale): """ Scale the parameters of a module and return it. """ for p in module.parameters(): p.detach().mul_(scale) return module def mean_flat(tensor): """ Take the mean over all non-batch dimensions. """ return tensor.mean(dim=list(range(1, len(tensor.shape)))) def normalization(channels): """ Make a standard normalization layer. :param channels: number of input channels. :return: an nn.Module for normalization. """ return GroupNorm32(32, channels) def timestep_embedding(timesteps, dim, max_period=10000): """ Create sinusoidal timestep embeddings. :param timesteps: a 1-D Tensor of N indices, one per batch element. These may be fractional. :param dim: the dimension of the output. :param max_period: controls the minimum frequency of the embeddings. :return: an [N x dim] Tensor of positional embeddings. """ half = dim // 2 freqs = th.exp( -math.log(max_period) * th.arange(start=0, end=half, dtype=th.float32) / half ).to(device=timesteps.device) args = timesteps[:, None].float() * freqs[None] embedding = th.cat([th.cos(args), th.sin(args)], dim=-1) if dim % 2: embedding = th.cat([embedding, th.zeros_like(embedding[:, :1])], dim=-1) return embedding def checkpoint(func, inputs, params, flag): """ Evaluate a function without caching intermediate activations, allowing for reduced memory at the expense of extra compute in the backward pass. :param func: the function to evaluate. :param inputs: the argument sequence to pass to `func`. :param params: a sequence of parameters `func` depends on but does not explicitly take as arguments. :param flag: if False, disable gradient checkpointing. """ if flag: args = tuple(inputs) + tuple(params) return CheckpointFunction.apply(func, len(inputs), *args) else: return func(*inputs) class CheckpointFunction(th.autograd.Function): @staticmethod def forward(ctx, run_function, length, *args): ctx.run_function = run_function ctx.input_tensors = list(args[:length]) ctx.input_params = list(args[length:]) with th.no_grad(): output_tensors = ctx.run_function(*ctx.input_tensors) return output_tensors @staticmethod def backward(ctx, *output_grads): ctx.input_tensors = [x.detach().requires_grad_(True) for x in ctx.input_tensors] with th.enable_grad(): # Fixes a bug where the first op in run_function modifies the # Tensor storage in place, which is not allowed for detach()'d # Tensors. shallow_copies = [x.view_as(x) for x in ctx.input_tensors] output_tensors = ctx.run_function(*shallow_copies) input_grads = th.autograd.grad( output_tensors, ctx.input_tensors + ctx.input_params, output_grads, allow_unused=True, ) del ctx.input_tensors del ctx.input_params del output_tensors return (None, None) + input_grads ================================================ FILE: guided_diffusion/resample.py ================================================ from abc import ABC, abstractmethod import numpy as np import torch as th import torch.distributed as dist def create_named_schedule_sampler(name, diffusion): """ Create a ScheduleSampler from a library of pre-defined samplers. :param name: the name of the sampler. :param diffusion: the diffusion object to sample for. """ if name == "uniform": return UniformSampler(diffusion) elif name == "loss-second-moment": return LossSecondMomentResampler(diffusion) else: raise NotImplementedError(f"unknown schedule sampler: {name}") class ScheduleSampler(ABC): """ A distribution over timesteps in the diffusion process, intended to reduce variance of the objective. By default, samplers perform unbiased importance sampling, in which the objective's mean is unchanged. However, subclasses may override sample() to change how the resampled terms are reweighted, allowing for actual changes in the objective. """ @abstractmethod def weights(self): """ Get a numpy array of weights, one per diffusion step. The weights needn't be normalized, but must be positive. """ def sample(self, batch_size, device): """ Importance-sample timesteps for a batch. :param batch_size: the number of timesteps. :param device: the torch device to save to. :return: a tuple (timesteps, weights): - timesteps: a tensor of timestep indices. - weights: a tensor of weights to scale the resulting losses. """ w = self.weights() p = w / np.sum(w) # loss별 weight를 normalize indices_np = np.random.choice(len(p), size=(batch_size,), p=p)[0] if indices_np == 0: indices_np += 1 assert indices_np > 0, f"indices_np : {indices_np}" indices_np = np.repeat(indices_np, batch_size) indices = th.from_numpy(indices_np).long().to(device) weights_np = 1 / (len(p) * p[indices_np]) weights = th.from_numpy(weights_np).float().to(device) return indices, weights class UniformSampler(ScheduleSampler): def __init__(self, diffusion): self.diffusion = diffusion self._weights = np.ones([diffusion.num_timesteps]) def weights(self): return self._weights class LossAwareSampler(ScheduleSampler): def update_with_local_losses(self, local_ts, local_losses): """ Update the reweighting using losses from a model. Call this method from each rank with a batch of timesteps and the corresponding losses for each of those timesteps. This method will perform synchronization to make sure all of the ranks maintain the exact same reweighting. :param local_ts: an integer Tensor of timesteps. :param local_losses: a 1D Tensor of losses. """ batch_sizes = [ th.tensor([0], dtype=th.int32, device=local_ts.device) for _ in range(dist.get_world_size()) ] dist.all_gather( batch_sizes, th.tensor([len(local_ts)], dtype=th.int32, device=local_ts.device), ) # Pad all_gather batches to be the maximum batch size. batch_sizes = [x.item() for x in batch_sizes] max_bs = max(batch_sizes) timestep_batches = [th.zeros(max_bs).to(local_ts) for bs in batch_sizes] loss_batches = [th.zeros(max_bs).to(local_losses) for bs in batch_sizes] dist.all_gather(timestep_batches, local_ts) dist.all_gather(loss_batches, local_losses) timesteps = [ x.item() for y, bs in zip(timestep_batches, batch_sizes) for x in y[:bs] ] losses = [x.item() for y, bs in zip(loss_batches, batch_sizes) for x in y[:bs]] self.update_with_all_losses(timesteps, losses) @abstractmethod def update_with_all_losses(self, ts, losses): """ Update the reweighting using losses from a model. Sub-classes should override this method to update the reweighting using losses from the model. This method directly updates the reweighting without synchronizing between workers. It is called by update_with_local_losses from all ranks with identical arguments. Thus, it should have deterministic behavior to maintain state across workers. :param ts: a list of int timesteps. :param losses: a list of float losses, one per timestep. """ class LossSecondMomentResampler(LossAwareSampler): def __init__(self, diffusion, history_per_term=10, uniform_prob=0.001): raise NotImplementedError self.diffusion = diffusion self.history_per_term = history_per_term self.uniform_prob = uniform_prob self._loss_history = np.zeros( [diffusion.num_timesteps, history_per_term], dtype=np.float64 ) self._loss_counts = np.zeros([diffusion.num_timesteps], dtype=np.int) def weights(self): if not self._warmed_up(): return np.ones([self.diffusion.num_timesteps], dtype=np.float64) weights = np.sqrt(np.mean(self._loss_history ** 2, axis=-1)) weights /= np.sum(weights) weights *= 1 - self.uniform_prob weights += self.uniform_prob / len(weights) return weights def update_with_all_losses(self, ts, losses): for t, loss in zip(ts, losses): if self._loss_counts[t] == self.history_per_term: # Shift out the oldest loss term. self._loss_history[t, :-1] = self._loss_history[t, 1:] self._loss_history[t, -1] = loss else: self._loss_history[t, self._loss_counts[t]] = loss self._loss_counts[t] += 1 def _warmed_up(self): return (self._loss_counts == self.history_per_term).all() ================================================ FILE: guided_diffusion/respace.py ================================================ import numpy as np import torch as th from .gaussian_diffusion import GaussianDiffusion def space_timesteps(num_timesteps, section_counts): """ Create a list of timesteps to use from an original diffusion process, given the number of timesteps we want to take from equally-sized portions of the original process. For example, if there's 300 timesteps and the section counts are [10,15,20] then the first 100 timesteps are strided to be 10 timesteps, the second 100 are strided to be 15 timesteps, and the final 100 are strided to be 20. If the stride is a string starting with "ddim", then the fixed striding from the DDIM paper is used, and only one section is allowed. :param num_timesteps: the number of diffusion steps in the original process to divide up. :param section_counts: either a list of numbers, or a string containing comma-separated numbers, indicating the step count per section. As a special case, use "ddimN" where N is a number of steps to use the striding from the DDIM paper. :return: a set of diffusion steps from the original process to use. """ if isinstance(section_counts, str): if section_counts.startswith("ddim"): desired_count = int(section_counts[len("ddim") :]) for i in range(1, num_timesteps): if len(range(0, num_timesteps, i)) == desired_count: return set(range(0, num_timesteps, i)) raise ValueError( f"cannot create exactly {num_timesteps} steps with an integer stride" ) section_counts = [int(x) for x in section_counts.split(",")] size_per = num_timesteps // len(section_counts) extra = num_timesteps % len(section_counts) start_idx = 0 all_steps = [] for i, section_count in enumerate(section_counts): size = size_per + (1 if i < extra else 0) if size < section_count: raise ValueError( f"cannot divide section of {size} steps into {section_count}" ) if section_count <= 1: frac_stride = 1 else: frac_stride = (size - 1) / (section_count - 1) cur_idx = 0.0 taken_steps = [] for _ in range(section_count): taken_steps.append(start_idx + round(cur_idx)) cur_idx += frac_stride all_steps += taken_steps start_idx += size return set(all_steps) class SpacedDiffusion(GaussianDiffusion): """ A diffusion process which can skip steps in a base diffusion process. :param use_timesteps: a collection (sequence or set) of timesteps from the original diffusion process to retain. :param kwargs: the kwargs to create the base diffusion process. """ def __init__(self, use_timesteps, **kwargs): self.use_timesteps = set(use_timesteps) self.timestep_map = [] self.original_num_steps = len(kwargs["betas"]) base_diffusion = GaussianDiffusion(**kwargs) # pylint: disable=missing-kwoa last_alpha_cumprod = 1.0 new_betas = [] for i, alpha_cumprod in enumerate(base_diffusion.alphas_cumprod): if i in self.use_timesteps: new_betas.append(1 - alpha_cumprod / last_alpha_cumprod) last_alpha_cumprod = alpha_cumprod self.timestep_map.append(i) kwargs["betas"] = np.array(new_betas) super().__init__(**kwargs) def p_mean_variance( self, model, *args, **kwargs ): # pylint: disable=signature-differs return super().p_mean_variance(self._wrap_model(model), *args, **kwargs) def training_losses( self, model, *args, **kwargs ): # pylint: disable=signature-differs return super().training_losses(self._wrap_model(model), *args, **kwargs) def condition_mean(self, cond_fn, *args, **kwargs): return super().condition_mean(self._wrap_model(cond_fn), *args, **kwargs) def condition_score(self, cond_fn, *args, **kwargs): return super().condition_score(self._wrap_model(cond_fn), *args, **kwargs) def _wrap_model(self, model): if isinstance(model, _WrappedModel): return model return _WrappedModel( model, self.timestep_map, self.rescale_timesteps, self.original_num_steps ) def _scale_timesteps(self, t): # Scaling is done by the wrapped model. return t class _WrappedModel: def __init__(self, model, timestep_map, rescale_timesteps, original_num_steps): self.model = model self.timestep_map = timestep_map self.rescale_timesteps = rescale_timesteps self.original_num_steps = original_num_steps def __call__(self, x, ts, **kwargs): map_tensor = th.tensor(self.timestep_map, device=ts.device, dtype=ts.dtype) new_ts = map_tensor[ts] if self.rescale_timesteps: new_ts = new_ts.float() * (1000.0 / self.original_num_steps) return self.model(x, new_ts, **kwargs) ================================================ FILE: guided_diffusion/script_util.py ================================================ import argparse import inspect from . import gaussian_diffusion as gd from .respace import SpacedDiffusion, space_timesteps from .unet import SuperResModel, UNetModel, EncoderUNetModel NUM_CLASSES = 1000 def diffusion_defaults(): """ Defaults for image and classifier training. """ return dict( learn_sigma=False, diffusion_steps=1000, noise_schedule="linear", timestep_respacing="", use_kl=False, predict_xstart=False, rescale_timesteps=False, rescale_learned_sigmas=False, ) def classifier_defaults(): """ Defaults for classifier models. """ return dict( image_size=64, classifier_use_fp16=False, classifier_width=128, classifier_depth=2, classifier_attention_resolutions="32,16,8", # 16 classifier_use_scale_shift_norm=True, # False classifier_resblock_updown=True, # False classifier_pool="attention", ) def model_and_diffusion_defaults(): """ Defaults for image training. """ res = dict( image_size=64, num_channels=128, num_res_blocks=2, num_heads=4, num_heads_upsample=-1, num_head_channels=-1, attention_resolutions="16,8", channel_mult="", dropout=0.0, class_cond=False, use_checkpoint=False, use_scale_shift_norm=True, resblock_updown=False, use_fp16=False, use_new_attention_order=False, ) res.update(diffusion_defaults()) return res def classifier_and_diffusion_defaults(): res = classifier_defaults() res.update(diffusion_defaults()) return res def create_model_and_diffusion( image_size, class_cond, learn_sigma, num_channels, num_res_blocks, channel_mult, num_heads, num_head_channels, num_heads_upsample, attention_resolutions, dropout, diffusion_steps, noise_schedule, timestep_respacing, use_kl, predict_xstart, rescale_timesteps, rescale_learned_sigmas, use_checkpoint, use_scale_shift_norm, resblock_updown, use_fp16, use_new_attention_order, ): model = create_model( image_size, num_channels, num_res_blocks, channel_mult=channel_mult, learn_sigma=learn_sigma, class_cond=class_cond, use_checkpoint=use_checkpoint, attention_resolutions=attention_resolutions, num_heads=num_heads, num_head_channels=num_head_channels, num_heads_upsample=num_heads_upsample, use_scale_shift_norm=use_scale_shift_norm, dropout=dropout, resblock_updown=resblock_updown, use_fp16=use_fp16, use_new_attention_order=use_new_attention_order, ) diffusion = create_gaussian_diffusion( steps=diffusion_steps, learn_sigma=learn_sigma, noise_schedule=noise_schedule, use_kl=use_kl, predict_xstart=predict_xstart, rescale_timesteps=rescale_timesteps, rescale_learned_sigmas=rescale_learned_sigmas, timestep_respacing=timestep_respacing, ) return model, diffusion def create_model( image_size, num_channels, num_res_blocks, channel_mult="", learn_sigma=False, class_cond=False, use_checkpoint=False, attention_resolutions="16", num_heads=1, num_head_channels=-1, num_heads_upsample=-1, use_scale_shift_norm=False, dropout=0, resblock_updown=False, use_fp16=False, use_new_attention_order=False, ): if channel_mult == "": if image_size == 512: channel_mult = (0.5, 1, 1, 2, 2, 4, 4) elif image_size == 256: channel_mult = (1, 1, 2, 2, 4, 4) elif image_size == 128: channel_mult = (1, 1, 2, 3, 4) elif image_size == 64: channel_mult = (1, 2, 3, 4) else: raise ValueError(f"unsupported image size: {image_size}") else: channel_mult = tuple(int(ch_mult) for ch_mult in channel_mult.split(",")) attention_ds = [] for res in attention_resolutions.split(","): attention_ds.append(image_size // int(res)) return UNetModel( image_size=image_size, in_channels=3, model_channels=num_channels, out_channels=(3 if not learn_sigma else 6), num_res_blocks=num_res_blocks, attention_resolutions=tuple(attention_ds), dropout=dropout, channel_mult=channel_mult, num_classes=(NUM_CLASSES if class_cond else None), use_checkpoint=use_checkpoint, use_fp16=use_fp16, num_heads=num_heads, num_head_channels=num_head_channels, num_heads_upsample=num_heads_upsample, use_scale_shift_norm=use_scale_shift_norm, resblock_updown=resblock_updown, use_new_attention_order=use_new_attention_order, ) def create_classifier_and_diffusion( image_size, classifier_use_fp16, classifier_width, classifier_depth, classifier_attention_resolutions, classifier_use_scale_shift_norm, classifier_resblock_updown, classifier_pool, learn_sigma, diffusion_steps, noise_schedule, timestep_respacing, use_kl, predict_xstart, rescale_timesteps, rescale_learned_sigmas, ): classifier = create_classifier( image_size, classifier_use_fp16, classifier_width, classifier_depth, classifier_attention_resolutions, classifier_use_scale_shift_norm, classifier_resblock_updown, classifier_pool, ) diffusion = create_gaussian_diffusion( steps=diffusion_steps, learn_sigma=learn_sigma, noise_schedule=noise_schedule, use_kl=use_kl, predict_xstart=predict_xstart, rescale_timesteps=rescale_timesteps, rescale_learned_sigmas=rescale_learned_sigmas, timestep_respacing=timestep_respacing, ) return classifier, diffusion def create_classifier( image_size, classifier_use_fp16, classifier_width, classifier_depth, classifier_attention_resolutions, classifier_use_scale_shift_norm, classifier_resblock_updown, classifier_pool, ): if image_size == 512: channel_mult = (0.5, 1, 1, 2, 2, 4, 4) elif image_size == 256: channel_mult = (1, 1, 2, 2, 4, 4) elif image_size == 128: channel_mult = (1, 1, 2, 3, 4) elif image_size == 64: channel_mult = (1, 2, 3, 4) else: raise ValueError(f"unsupported image size: {image_size}") attention_ds = [] for res in classifier_attention_resolutions.split(","): attention_ds.append(image_size // int(res)) return EncoderUNetModel( image_size=image_size, in_channels=3, model_channels=classifier_width, out_channels=1000, num_res_blocks=classifier_depth, attention_resolutions=tuple(attention_ds), channel_mult=channel_mult, use_fp16=classifier_use_fp16, num_head_channels=64, use_scale_shift_norm=classifier_use_scale_shift_norm, resblock_updown=classifier_resblock_updown, pool=classifier_pool, ) def sr_model_and_diffusion_defaults(): res = model_and_diffusion_defaults() res["large_size"] = 256 res["small_size"] = 64 arg_names = inspect.getfullargspec(sr_create_model_and_diffusion)[0] for k in res.copy().keys(): if k not in arg_names: del res[k] return res def sr_create_model_and_diffusion( large_size, small_size, class_cond, learn_sigma, num_channels, num_res_blocks, num_heads, num_head_channels, num_heads_upsample, attention_resolutions, dropout, diffusion_steps, noise_schedule, timestep_respacing, use_kl, predict_xstart, rescale_timesteps, rescale_learned_sigmas, use_checkpoint, use_scale_shift_norm, resblock_updown, use_fp16, ): model = sr_create_model( large_size, small_size, num_channels, num_res_blocks, learn_sigma=learn_sigma, class_cond=class_cond, use_checkpoint=use_checkpoint, attention_resolutions=attention_resolutions, num_heads=num_heads, num_head_channels=num_head_channels, num_heads_upsample=num_heads_upsample, use_scale_shift_norm=use_scale_shift_norm, dropout=dropout, resblock_updown=resblock_updown, use_fp16=use_fp16, ) diffusion = create_gaussian_diffusion( steps=diffusion_steps, learn_sigma=learn_sigma, noise_schedule=noise_schedule, use_kl=use_kl, predict_xstart=predict_xstart, rescale_timesteps=rescale_timesteps, rescale_learned_sigmas=rescale_learned_sigmas, timestep_respacing=timestep_respacing, ) return model, diffusion def sr_create_model( large_size, small_size, num_channels, num_res_blocks, learn_sigma, class_cond, use_checkpoint, attention_resolutions, num_heads, num_head_channels, num_heads_upsample, use_scale_shift_norm, dropout, resblock_updown, use_fp16, ): _ = small_size # hack to prevent unused variable if large_size == 512: channel_mult = (1, 1, 2, 2, 4, 4) elif large_size == 256: channel_mult = (1, 1, 2, 2, 4, 4) elif large_size == 64: channel_mult = (1, 2, 3, 4) else: raise ValueError(f"unsupported large size: {large_size}") attention_ds = [] for res in attention_resolutions.split(","): attention_ds.append(large_size // int(res)) return SuperResModel( image_size=large_size, in_channels=3, model_channels=num_channels, out_channels=(3 if not learn_sigma else 6), num_res_blocks=num_res_blocks, attention_resolutions=tuple(attention_ds), dropout=dropout, channel_mult=channel_mult, num_classes=(NUM_CLASSES if class_cond else None), use_checkpoint=use_checkpoint, num_heads=num_heads, num_head_channels=num_head_channels, num_heads_upsample=num_heads_upsample, use_scale_shift_norm=use_scale_shift_norm, resblock_updown=resblock_updown, use_fp16=use_fp16, ) def create_gaussian_diffusion( *, steps=1000, learn_sigma=False, sigma_small=False, noise_schedule="linear", use_kl=False, predict_xstart=False, rescale_timesteps=False, rescale_learned_sigmas=False, timestep_respacing="", ): betas = gd.get_named_beta_schedule(noise_schedule, steps) if use_kl: loss_type = gd.LossType.RESCALED_KL elif rescale_learned_sigmas: loss_type = gd.LossType.RESCALED_MSE else: loss_type = gd.LossType.MSE if not timestep_respacing: timestep_respacing = [steps] return SpacedDiffusion( use_timesteps=space_timesteps(steps, timestep_respacing), betas=betas, model_mean_type=( gd.ModelMeanType.EPSILON if not predict_xstart else gd.ModelMeanType.START_X ), model_var_type=( ( gd.ModelVarType.FIXED_LARGE if not sigma_small else gd.ModelVarType.FIXED_SMALL ) if not learn_sigma else gd.ModelVarType.LEARNED_RANGE ), loss_type=loss_type, rescale_timesteps=rescale_timesteps, ) def add_dict_to_argparser(parser, default_dict): for k, v in default_dict.items(): v_type = type(v) if v is None: v_type = str elif isinstance(v, bool): v_type = str2bool parser.add_argument(f"--{k}", default=v, type=v_type) def args_to_dict(args, keys): return {k: getattr(args, k) for k in keys} def str2bool(v): """ https://stackoverflow.com/questions/15008758/parsing-boolean-values-with-argparse """ if isinstance(v, bool): return v if v.lower() in ("yes", "true", "t", "y", "1"): return True elif v.lower() in ("no", "false", "f", "n", "0"): return False else: raise argparse.ArgumentTypeError("boolean value expected") ================================================ FILE: guided_diffusion/train_util.py ================================================ import copy import functools import os import blobfile as bf import torch as th import torch.distributed as dist from torch.nn.parallel.distributed import DistributedDataParallel as DDP from torch.optim import AdamW from . import dist_util, logger from .fp16_util import MixedPrecisionTrainer from .nn import update_ema from .resample import LossAwareSampler, UniformSampler # For ImageNet experiments, this was a good default value. # We found that the lg_loss_scale quickly climbed to # 20-21 within the first ~1K steps of training. INITIAL_LOG_LOSS_SCALE = 20.0 class TrainLoop: def __init__( self, *, model, diffusion, data, batch_size, microbatch, lr, ema_rate, log_interval, save_interval, resume_checkpoint, use_fp16=False, fp16_scale_growth=1e-3, schedule_sampler=None, weight_decay=0.0, lr_anneal_steps=0, ): self.model = model self.diffusion = diffusion self.data = data self.batch_size = batch_size self.microbatch = microbatch if microbatch > 0 else batch_size self.lr = lr self.ema_rate = ( [ema_rate] if isinstance(ema_rate, float) else [float(x) for x in ema_rate.split(",")] ) self.log_interval = log_interval self.save_interval = save_interval self.resume_checkpoint = resume_checkpoint self.use_fp16 = use_fp16 self.fp16_scale_growth = fp16_scale_growth self.schedule_sampler = schedule_sampler or UniformSampler(diffusion) self.weight_decay = weight_decay self.lr_anneal_steps = lr_anneal_steps self.step = 0 self.resume_step = 0 self.global_batch = self.batch_size * dist.get_world_size() self.sync_cuda = th.cuda.is_available() self._load_and_sync_parameters() self.mp_trainer = MixedPrecisionTrainer( model=self.model, use_fp16=self.use_fp16, fp16_scale_growth=fp16_scale_growth, ) self.opt = AdamW( self.mp_trainer.master_params, lr=self.lr, weight_decay=self.weight_decay ) if self.resume_step: self._load_optimizer_state() # Model was resumed, either due to a restart or a checkpoint # being specified at the command line. self.ema_params = [ self._load_ema_parameters(rate) for rate in self.ema_rate ] else: self.ema_params = [ copy.deepcopy(self.mp_trainer.master_params) for _ in range(len(self.ema_rate)) ] if th.cuda.is_available(): self.use_ddp = True self.ddp_model = DDP( self.model, device_ids=[dist_util.dev()], output_device=dist_util.dev(), broadcast_buffers=False, bucket_cap_mb=128, find_unused_parameters=False, ) else: if dist.get_world_size() > 1: logger.warn( "Distributed training requires CUDA. " "Gradients will not be synchronized properly!" ) self.use_ddp = False self.ddp_model = self.model def _load_and_sync_parameters(self): resume_checkpoint = find_resume_checkpoint() or self.resume_checkpoint if resume_checkpoint: self.resume_step = parse_resume_step_from_filename(resume_checkpoint) if dist.get_rank() == 0: logger.log(f"loading model from checkpoint: {resume_checkpoint}...") self.model.load_state_dict( dist_util.load_state_dict( resume_checkpoint, map_location=dist_util.dev() ) ) dist_util.sync_params(self.model.parameters()) def _load_ema_parameters(self, rate): ema_params = copy.deepcopy(self.mp_trainer.master_params) main_checkpoint = find_resume_checkpoint() or self.resume_checkpoint ema_checkpoint = find_ema_checkpoint(main_checkpoint, self.resume_step, rate) if ema_checkpoint: if dist.get_rank() == 0: logger.log(f"loading EMA from checkpoint: {ema_checkpoint}...") state_dict = dist_util.load_state_dict( ema_checkpoint, map_location=dist_util.dev() ) ema_params = self.mp_trainer.state_dict_to_master_params(state_dict) dist_util.sync_params(ema_params) return ema_params def _load_optimizer_state(self): main_checkpoint = find_resume_checkpoint() or self.resume_checkpoint opt_checkpoint = bf.join( bf.dirname(main_checkpoint), f"opt{self.resume_step:06}.pt" ) if bf.exists(opt_checkpoint): logger.log(f"loading optimizer state from checkpoint: {opt_checkpoint}") state_dict = dist_util.load_state_dict( opt_checkpoint, map_location=dist_util.dev() ) self.opt.load_state_dict(state_dict) def run_loop(self): while ( not self.lr_anneal_steps or self.step + self.resume_step < self.lr_anneal_steps ): batch, cond = next(self.data) self.run_step(batch, cond) if self.step % self.log_interval == 0: logger.dumpkvs() if self.step % self.save_interval == 0: self.save() # Run for a finite amount of time in integration tests. if os.environ.get("DIFFUSION_TRAINING_TEST", "") and self.step > 0: return self.step += 1 # Save the last checkpoint if it wasn't already saved. if (self.step - 1) % self.save_interval != 0: self.save() def run_step(self, batch, cond): self.forward_backward(batch, cond) took_step = self.mp_trainer.optimize(self.opt) if took_step: self._update_ema() self._anneal_lr() self.log_step() def forward_backward(self, batch, cond): self.mp_trainer.zero_grad() for i in range(0, batch.shape[0], self.microbatch): micro = batch[i : i + self.microbatch].to(dist_util.dev()) micro_cond = { k: v[i : i + self.microbatch].to(dist_util.dev()) for k, v in cond.items() } last_batch = (i + self.microbatch) >= batch.shape[0] t, weights = self.schedule_sampler.sample(micro.shape[0], dist_util.dev()) # print("TrainLoop/forward_backward() t:", t) assert t[0] >= 0, f"train_util.py t[0] : {t[0]}" compute_losses = functools.partial( self.diffusion.training_losses, self.ddp_model, micro, t, model_kwargs=micro_cond, ) if last_batch or not self.use_ddp: losses = compute_losses() else: with self.ddp_model.no_sync(): losses = compute_losses() if isinstance(self.schedule_sampler, LossAwareSampler): self.schedule_sampler.update_with_local_losses( t, losses["loss"].detach() ) loss = (losses["loss"] * weights).mean() log_loss_dict( self.diffusion, t, {k: v * weights for k, v in losses.items()} ) self.mp_trainer.backward(loss) def _update_ema(self): for rate, params in zip(self.ema_rate, self.ema_params): update_ema(params, self.mp_trainer.master_params, rate=rate) def _anneal_lr(self): if not self.lr_anneal_steps: return frac_done = (self.step + self.resume_step) / self.lr_anneal_steps lr = self.lr * (1 - frac_done) for param_group in self.opt.param_groups: param_group["lr"] = lr def log_step(self): logger.logkv("step", self.step + self.resume_step) logger.logkv("samples", (self.step + self.resume_step + 1) * self.global_batch) def save(self): def save_checkpoint(rate, params): state_dict = self.mp_trainer.master_params_to_state_dict(params) if dist.get_rank() == 0: logger.log(f"saving model {rate}...") if not rate: filename = f"model{(self.step+self.resume_step):06d}.pt" else: filename = f"ema_{rate}_{(self.step+self.resume_step):06d}.pt" with bf.BlobFile(bf.join(get_blob_logdir(), filename), "wb") as f: th.save(state_dict, f) save_checkpoint(0, self.mp_trainer.master_params) for rate, params in zip(self.ema_rate, self.ema_params): save_checkpoint(rate, params) if dist.get_rank() == 0: with bf.BlobFile( bf.join(get_blob_logdir(), f"opt{(self.step+self.resume_step):06d}.pt"), "wb", ) as f: th.save(self.opt.state_dict(), f) dist.barrier() def parse_resume_step_from_filename(filename): """ Parse filenames of the form path/to/modelNNNNNN.pt, where NNNNNN is the checkpoint's number of steps. """ split = filename.split("model") if len(split) < 2: return 0 split1 = split[-1].split(".")[0] try: return int(split1) except ValueError: return 0 def get_blob_logdir(): # You can change this to be a separate path to save checkpoints to # a blobstore or some external drive. return logger.get_dir() def find_resume_checkpoint(): # On your infrastructure, you may want to override this to automatically # discover the latest checkpoint on your blob storage, etc. return None def find_ema_checkpoint(main_checkpoint, step, rate): if main_checkpoint is None: return None filename = f"ema_{rate}_{(step):06d}.pt" path = bf.join(bf.dirname(main_checkpoint), filename) if bf.exists(path): return path return None def log_loss_dict(diffusion, ts, losses): for key, values in losses.items(): logger.logkv_mean(key, values.mean().item()) # Log the quantiles (four quartiles, in particular). for sub_t, sub_loss in zip(ts.cpu().numpy(), values.detach().cpu().numpy()): quartile = int(4 * sub_t / diffusion.num_timesteps) logger.logkv_mean(f"{key}_q{quartile}", sub_loss) ================================================ FILE: guided_diffusion/unet.py ================================================ from abc import abstractmethod import math import numpy as np import torch as th import torch.nn as nn import torch.nn.functional as F from .fp16_util import convert_module_to_f16, convert_module_to_f32 from .nn import ( checkpoint, conv_nd, linear, avg_pool_nd, zero_module, normalization, timestep_embedding, ) class AttentionPool2d(nn.Module): """ Adapted from CLIP: https://github.com/openai/CLIP/blob/main/clip/model.py """ def __init__( self, spacial_dim: int, embed_dim: int, num_heads_channels: int, output_dim: int = None, ): super().__init__() self.positional_embedding = nn.Parameter( th.randn(embed_dim, spacial_dim ** 2 + 1) / embed_dim ** 0.5 ) self.qkv_proj = conv_nd(1, embed_dim, 3 * embed_dim, 1) self.c_proj = conv_nd(1, embed_dim, output_dim or embed_dim, 1) self.num_heads = embed_dim // num_heads_channels self.attention = QKVAttention(self.num_heads) def forward(self, x): b, c, *_spatial = x.shape x = x.reshape(b, c, -1) # NC(HW) x = th.cat([x.mean(dim=-1, keepdim=True), x], dim=-1) # NC(HW+1) x = x + self.positional_embedding[None, :, :].to(x.dtype) # NC(HW+1) x = self.qkv_proj(x) x = self.attention(x) x = self.c_proj(x) return x[:, :, 0] class TimestepBlock(nn.Module): """ Any module where forward() takes timestep embeddings as a second argument. """ @abstractmethod def forward(self, x, emb): """ Apply the module to `x` given `emb` timestep embeddings. """ class TimestepEmbedSequential(nn.Sequential, TimestepBlock): """ A sequential module that passes timestep embeddings to the children that support it as an extra input. """ def forward(self, x, emb): for layer in self: if isinstance(layer, TimestepBlock): x = layer(x, emb) else: x = layer(x) return x class Upsample(nn.Module): """ An upsampling layer with an optional convolution. :param channels: channels in the inputs and outputs. :param use_conv: a bool determining if a convolution is applied. :param dims: determines if the signal is 1D, 2D, or 3D. If 3D, then upsampling occurs in the inner-two dimensions. """ def __init__(self, channels, use_conv, dims=2, out_channels=None): super().__init__() self.channels = channels self.out_channels = out_channels or channels self.use_conv = use_conv self.dims = dims if use_conv: self.conv = conv_nd(dims, self.channels, self.out_channels, 3, padding=1) def forward(self, x): assert x.shape[1] == self.channels if self.dims == 3: x = F.interpolate( x, (x.shape[2], x.shape[3] * 2, x.shape[4] * 2), mode="nearest" ) else: x = F.interpolate(x, scale_factor=2, mode="nearest") if self.use_conv: x = self.conv(x) return x class Downsample(nn.Module): """ A downsampling layer with an optional convolution. :param channels: channels in the inputs and outputs. :param use_conv: a bool determining if a convolution is applied. :param dims: determines if the signal is 1D, 2D, or 3D. If 3D, then downsampling occurs in the inner-two dimensions. """ def __init__(self, channels, use_conv, dims=2, out_channels=None): super().__init__() self.channels = channels self.out_channels = out_channels or channels self.use_conv = use_conv self.dims = dims stride = 2 if dims != 3 else (1, 2, 2) if use_conv: self.op = conv_nd( dims, self.channels, self.out_channels, 3, stride=stride, padding=1 ) else: assert self.channels == self.out_channels self.op = avg_pool_nd(dims, kernel_size=stride, stride=stride) def forward(self, x): assert x.shape[1] == self.channels return self.op(x) class ResBlock(TimestepBlock): """ A residual block that can optionally change the number of channels. :param channels: the number of input channels. :param emb_channels: the number of timestep embedding channels. :param dropout: the rate of dropout. :param out_channels: if specified, the number of out channels. :param use_conv: if True and out_channels is specified, use a spatial convolution instead of a smaller 1x1 convolution to change the channels in the skip connection. :param dims: determines if the signal is 1D, 2D, or 3D. :param use_checkpoint: if True, use gradient checkpointing on this module. :param up: if True, use this block for upsampling. :param down: if True, use this block for downsampling. """ def __init__( self, channels, emb_channels, dropout, out_channels=None, use_conv=False, use_scale_shift_norm=False, dims=2, use_checkpoint=False, up=False, down=False, ): super().__init__() self.channels = channels self.emb_channels = emb_channels self.dropout = dropout self.out_channels = out_channels or channels self.use_conv = use_conv self.use_checkpoint = use_checkpoint self.use_scale_shift_norm = use_scale_shift_norm self.in_layers = nn.Sequential( normalization(channels), nn.SiLU(), conv_nd(dims, channels, self.out_channels, 3, padding=1), ) self.updown = up or down if up: self.h_upd = Upsample(channels, False, dims) self.x_upd = Upsample(channels, False, dims) elif down: self.h_upd = Downsample(channels, False, dims) self.x_upd = Downsample(channels, False, dims) else: self.h_upd = self.x_upd = nn.Identity() self.emb_layers = nn.Sequential( nn.SiLU(), linear( emb_channels, 2 * self.out_channels if use_scale_shift_norm else self.out_channels, ), ) self.out_layers = nn.Sequential( normalization(self.out_channels), nn.SiLU(), nn.Dropout(p=dropout), zero_module( conv_nd(dims, self.out_channels, self.out_channels, 3, padding=1) ), ) if self.out_channels == channels: self.skip_connection = nn.Identity() elif use_conv: self.skip_connection = conv_nd( dims, channels, self.out_channels, 3, padding=1 ) else: self.skip_connection = conv_nd(dims, channels, self.out_channels, 1) def forward(self, x, emb): """ Apply the block to a Tensor, conditioned on a timestep embedding. :param x: an [N x C x ...] Tensor of features. :param emb: an [N x emb_channels] Tensor of timestep embeddings. :return: an [N x C x ...] Tensor of outputs. """ return checkpoint( self._forward, (x, emb), self.parameters(), self.use_checkpoint ) def _forward(self, x, emb): if self.updown: in_rest, in_conv = self.in_layers[:-1], self.in_layers[-1] h = in_rest(x) h = self.h_upd(h) x = self.x_upd(x) h = in_conv(h) else: h = self.in_layers(x) emb_out = self.emb_layers(emb).type(h.dtype) while len(emb_out.shape) < len(h.shape): emb_out = emb_out[..., None] if self.use_scale_shift_norm: out_norm, out_rest = self.out_layers[0], self.out_layers[1:] scale, shift = th.chunk(emb_out, 2, dim=1) h = out_norm(h) * (1 + scale) + shift h = out_rest(h) else: h = h + emb_out h = self.out_layers(h) return self.skip_connection(x) + h class AttentionBlock(nn.Module): """ An attention block that allows spatial positions to attend to each other. Originally ported from here, but adapted to the N-d case. https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/models/unet.py#L66. """ def __init__( self, channels, num_heads=1, num_head_channels=-1, use_checkpoint=False, use_new_attention_order=False, ): super().__init__() self.channels = channels if num_head_channels == -1: self.num_heads = num_heads else: assert ( channels % num_head_channels == 0 ), f"q,k,v channels {channels} is not divisible by num_head_channels {num_head_channels}" self.num_heads = channels // num_head_channels self.use_checkpoint = use_checkpoint self.norm = normalization(channels) self.qkv = conv_nd(1, channels, channels * 3, 1) if use_new_attention_order: # split qkv before split heads self.attention = QKVAttention(self.num_heads) else: # split heads before split qkv self.attention = QKVAttentionLegacy(self.num_heads) self.proj_out = zero_module(conv_nd(1, channels, channels, 1)) def forward(self, x): return checkpoint(self._forward, (x,), self.parameters(), True) def _forward(self, x): b, c, *spatial = x.shape x = x.reshape(b, c, -1) qkv = self.qkv(self.norm(x)) h = self.attention(qkv) h = self.proj_out(h) return (x + h).reshape(b, c, *spatial) def count_flops_attn(model, _x, y): """ A counter for the `thop` package to count the operations in an attention operation. Meant to be used like: macs, params = thop.profile( model, inputs=(inputs, timestamps), custom_ops={QKVAttention: QKVAttention.count_flops}, ) """ b, c, *spatial = y[0].shape num_spatial = int(np.prod(spatial)) # We perform two matmuls with the same number of ops. # The first computes the weight matrix, the second computes # the combination of the value vectors. matmul_ops = 2 * b * (num_spatial ** 2) * c model.total_ops += th.DoubleTensor([matmul_ops]) class QKVAttentionLegacy(nn.Module): """ A module which performs QKV attention. Matches legacy QKVAttention + input/ouput heads shaping """ def __init__(self, n_heads): super().__init__() self.n_heads = n_heads def forward(self, qkv): """ Apply QKV attention. :param qkv: an [N x (H * 3 * C) x T] tensor of Qs, Ks, and Vs. :return: an [N x (H * C) x T] tensor after attention. """ bs, width, length = qkv.shape assert width % (3 * self.n_heads) == 0 ch = width // (3 * self.n_heads) q, k, v = qkv.reshape(bs * self.n_heads, ch * 3, length).split(ch, dim=1) scale = 1 / math.sqrt(math.sqrt(ch)) weight = th.einsum( "bct,bcs->bts", q * scale, k * scale ) # More stable with f16 than dividing afterwards weight = th.softmax(weight.float(), dim=-1).type(weight.dtype) a = th.einsum("bts,bcs->bct", weight, v) return a.reshape(bs, -1, length) @staticmethod def count_flops(model, _x, y): return count_flops_attn(model, _x, y) class QKVAttention(nn.Module): """ A module which performs QKV attention and splits in a different order. """ def __init__(self, n_heads): super().__init__() self.n_heads = n_heads def forward(self, qkv): """ Apply QKV attention. :param qkv: an [N x (3 * H * C) x T] tensor of Qs, Ks, and Vs. :return: an [N x (H * C) x T] tensor after attention. """ bs, width, length = qkv.shape assert width % (3 * self.n_heads) == 0 ch = width // (3 * self.n_heads) q, k, v = qkv.chunk(3, dim=1) scale = 1 / math.sqrt(math.sqrt(ch)) weight = th.einsum( "bct,bcs->bts", (q * scale).view(bs * self.n_heads, ch, length), (k * scale).view(bs * self.n_heads, ch, length), ) # More stable with f16 than dividing afterwards weight = th.softmax(weight.float(), dim=-1).type(weight.dtype) a = th.einsum("bts,bcs->bct", weight, v.reshape(bs * self.n_heads, ch, length)) return a.reshape(bs, -1, length) @staticmethod def count_flops(model, _x, y): return count_flops_attn(model, _x, y) class UNetModel(nn.Module): """ The full UNet model with attention and timestep embedding. :param in_channels: channels in the input Tensor. :param model_channels: base channel count for the model. :param out_channels: channels in the output Tensor. :param num_res_blocks: number of residual blocks per downsample. :param attention_resolutions: a collection of downsample rates at which attention will take place. May be a set, list, or tuple. For example, if this contains 4, then at 4x downsampling, attention will be used. :param dropout: the dropout probability. :param channel_mult: channel multiplier for each level of the UNet. :param conv_resample: if True, use learned convolutions for upsampling and downsampling. :param dims: determines if the signal is 1D, 2D, or 3D. :param num_classes: if specified (as an int), then this model will be class-conditional with `num_classes` classes. :param use_checkpoint: use gradient checkpointing to reduce memory usage. :param num_heads: the number of attention heads in each attention layer. :param num_heads_channels: if specified, ignore num_heads and instead use a fixed channel width per attention head. :param num_heads_upsample: works with num_heads to set a different number of heads for upsampling. Deprecated. :param use_scale_shift_norm: use a FiLM-like conditioning mechanism. :param resblock_updown: use residual blocks for up/downsampling. :param use_new_attention_order: use a different attention pattern for potentially increased efficiency. """ def __init__( self, image_size, in_channels, model_channels, out_channels, blur, num_res_blocks = 3, attention_resolutions = [4,8], dropout=0, channel_mult=(1, 2, 3, 4), conv_resample=True, dims=2, num_classes=None, use_checkpoint=False, use_fp16=False, num_heads=4, num_head_channels=-1, num_heads_upsample=-1, use_scale_shift_norm=True, resblock_updown=False, use_new_attention_order=False, freq_feat = False ): super().__init__() if num_heads_upsample == -1: num_heads_upsample = num_heads self.image_size = image_size self.in_channels = in_channels self.model_channels = model_channels self.out_channels = out_channels self.num_res_blocks = num_res_blocks self.attention_resolutions = attention_resolutions self.dropout = dropout self.channel_mult = channel_mult self.conv_resample = conv_resample self.num_classes = num_classes self.use_checkpoint = use_checkpoint self.dtype = th.float16 if use_fp16 else th.float32 self.num_heads = num_heads self.num_head_channels = num_head_channels self.num_heads_upsample = num_heads_upsample self.freq_feat = freq_feat self.blur = blur time_embed_dim = model_channels * 4 self.time_embed = nn.Sequential( linear(model_channels, time_embed_dim), nn.SiLU(), linear(time_embed_dim, time_embed_dim), ) if freq_feat: self.in_channels += 3 in_channels += 3 if self.num_classes is not None: self.label_emb = nn.Embedding(num_classes, time_embed_dim) ch = input_ch = int(channel_mult[0] * model_channels) self.input_blocks = nn.ModuleList( [TimestepEmbedSequential(conv_nd(dims, in_channels, ch, 3, padding=1))] ) print("input_blocks", self.input_blocks[0][0].weight.shape) if freq_feat: # Set the weights to the frequency features to zeros self.input_blocks[0][0].weight.data[:, 3: :, :] = 0 # for j in range(6): # w = self.input_blocks[0][0].weight[:, j, :, :] # print(f"j: {j} w.mean(): {w.mean()} w.std(): {w.std()}") # raise NotImplementedError self._feature_size = ch input_block_chans = [ch] ds = 1 for level, mult in enumerate(channel_mult): for _ in range(num_res_blocks): layers = [ ResBlock( ch, time_embed_dim, dropout, out_channels=int(mult * model_channels), dims=dims, use_checkpoint=use_checkpoint, use_scale_shift_norm=use_scale_shift_norm, ) ] ch = int(mult * model_channels) if ds in attention_resolutions: layers.append( AttentionBlock( ch, use_checkpoint=use_checkpoint, num_heads=num_heads, num_head_channels=num_head_channels, use_new_attention_order=use_new_attention_order, ) ) self.input_blocks.append(TimestepEmbedSequential(*layers)) self._feature_size += ch input_block_chans.append(ch) if level != len(channel_mult) - 1: out_ch = ch self.input_blocks.append( TimestepEmbedSequential( ResBlock( ch, time_embed_dim, dropout, out_channels=out_ch, dims=dims, use_checkpoint=use_checkpoint, use_scale_shift_norm=use_scale_shift_norm, down=True, ) if resblock_updown else Downsample( ch, conv_resample, dims=dims, out_channels=out_ch ) ) ) ch = out_ch input_block_chans.append(ch) ds *= 2 self._feature_size += ch self.middle_block = TimestepEmbedSequential( ResBlock( ch, time_embed_dim, dropout, dims=dims, use_checkpoint=use_checkpoint, use_scale_shift_norm=use_scale_shift_norm, ), AttentionBlock( ch, use_checkpoint=use_checkpoint, num_heads=num_heads, num_head_channels=num_head_channels, use_new_attention_order=use_new_attention_order, ), ResBlock( ch, time_embed_dim, dropout, dims=dims, use_checkpoint=use_checkpoint, use_scale_shift_norm=use_scale_shift_norm, ), ) self._feature_size += ch self.output_blocks = nn.ModuleList([]) for level, mult in list(enumerate(channel_mult))[::-1]: for i in range(num_res_blocks + 1): ich = input_block_chans.pop() layers = [ ResBlock( ch + ich, time_embed_dim, dropout, out_channels=int(model_channels * mult), dims=dims, use_checkpoint=use_checkpoint, use_scale_shift_norm=use_scale_shift_norm, ) ] ch = int(model_channels * mult) if ds in attention_resolutions: layers.append( AttentionBlock( ch, use_checkpoint=use_checkpoint, num_heads=num_heads_upsample, num_head_channels=num_head_channels, use_new_attention_order=use_new_attention_order, ) ) if level and i == num_res_blocks: out_ch = ch layers.append( ResBlock( ch, time_embed_dim, dropout, out_channels=out_ch, dims=dims, use_checkpoint=use_checkpoint, use_scale_shift_norm=use_scale_shift_norm, up=True, ) if resblock_updown else Upsample(ch, conv_resample, dims=dims, out_channels=out_ch) ) ds //= 2 self.output_blocks.append(TimestepEmbedSequential(*layers)) self._feature_size += ch self.out = nn.Sequential( normalization(ch), nn.SiLU(), zero_module(conv_nd(dims, input_ch, out_channels, 3, padding=1)), ) def convert_to_fp16(self): """ Convert the torso of the model to float16. """ self.input_blocks.apply(convert_module_to_f16) self.middle_block.apply(convert_module_to_f16) self.output_blocks.apply(convert_module_to_f16) def convert_to_fp32(self): """ Convert the torso of the model to float32. """ self.input_blocks.apply(convert_module_to_f32) self.middle_block.apply(convert_module_to_f32) self.output_blocks.apply(convert_module_to_f32) def forward(self, x, timesteps, y=None): """ Apply the model to an input batch. :param x: an [N x C x ...] Tensor of inputs. :param timesteps: a 1-D batch of timesteps. :param y: an [N] Tensor of labels, if class-conditional. :return: an [N x C x ...] Tensor of outputs. """ assert (y is not None) == ( self.num_classes is not None ), "must specify y if and only if the model is class-conditional" # timesteps = [timesteps] * x.shape[0] # timesteps = th.tensor(timesteps).to(x.device) if self.freq_feat: x_freq = self.blur.Ut(x).view(x.shape) x = th.cat([x, x_freq], dim=1) hs = [] emb = self.time_embed(timestep_embedding(timesteps, self.model_channels)) if self.num_classes is not None: assert y.shape == (x.shape[0],) emb = emb + self.label_emb(y) h = x.type(self.dtype) for module in self.input_blocks: h = module(h, emb) hs.append(h) h = self.middle_block(h, emb) for module in self.output_blocks: h = th.cat([h, hs.pop()], dim=1) h = module(h, emb) h = h.type(x.dtype) return self.out(h) class SuperResModel(UNetModel): """ A UNetModel that performs super-resolution. Expects an extra kwarg `low_res` to condition on a low-resolution image. """ def __init__(self, image_size, in_channels, *args, **kwargs): super().__init__(image_size, in_channels * 2, *args, **kwargs) def forward(self, x, timesteps, low_res=None, **kwargs): _, _, new_height, new_width = x.shape upsampled = F.interpolate(low_res, (new_height, new_width), mode="bilinear") x = th.cat([x, upsampled], dim=1) return super().forward(x, timesteps, **kwargs) class EncoderUNetModel(nn.Module): """ The half UNet model with attention and timestep embedding. For usage, see UNet. """ def __init__( self, image_size, in_channels, model_channels, out_channels, num_res_blocks, attention_resolutions, dropout=0, channel_mult=(1, 2, 4, 8), conv_resample=True, dims=2, use_checkpoint=False, use_fp16=False, num_heads=1, num_head_channels=-1, num_heads_upsample=-1, use_scale_shift_norm=False, resblock_updown=False, use_new_attention_order=False, pool="adaptive", ): super().__init__() if num_heads_upsample == -1: num_heads_upsample = num_heads self.in_channels = in_channels self.model_channels = model_channels self.out_channels = out_channels self.num_res_blocks = num_res_blocks self.attention_resolutions = attention_resolutions self.dropout = dropout self.channel_mult = channel_mult self.conv_resample = conv_resample self.use_checkpoint = use_checkpoint self.dtype = th.float16 if use_fp16 else th.float32 self.num_heads = num_heads self.num_head_channels = num_head_channels self.num_heads_upsample = num_heads_upsample time_embed_dim = model_channels * 4 self.time_embed = nn.Sequential( linear(model_channels, time_embed_dim), nn.SiLU(), linear(time_embed_dim, time_embed_dim), ) ch = int(channel_mult[0] * model_channels) self.input_blocks = nn.ModuleList( [TimestepEmbedSequential(conv_nd(dims, in_channels, ch, 3, padding=1))] ) self._feature_size = ch input_block_chans = [ch] ds = 1 for level, mult in enumerate(channel_mult): for _ in range(num_res_blocks): layers = [ ResBlock( ch, time_embed_dim, dropout, out_channels=int(mult * model_channels), dims=dims, use_checkpoint=use_checkpoint, use_scale_shift_norm=use_scale_shift_norm, ) ] ch = int(mult * model_channels) if ds in attention_resolutions: layers.append( AttentionBlock( ch, use_checkpoint=use_checkpoint, num_heads=num_heads, num_head_channels=num_head_channels, use_new_attention_order=use_new_attention_order, ) ) self.input_blocks.append(TimestepEmbedSequential(*layers)) self._feature_size += ch input_block_chans.append(ch) if level != len(channel_mult) - 1: out_ch = ch self.input_blocks.append( TimestepEmbedSequential( ResBlock( ch, time_embed_dim, dropout, out_channels=out_ch, dims=dims, use_checkpoint=use_checkpoint, use_scale_shift_norm=use_scale_shift_norm, down=True, ) if resblock_updown else Downsample( ch, conv_resample, dims=dims, out_channels=out_ch ) ) ) ch = out_ch input_block_chans.append(ch) ds *= 2 self._feature_size += ch self.middle_block = TimestepEmbedSequential( ResBlock( ch, time_embed_dim, dropout, dims=dims, use_checkpoint=use_checkpoint, use_scale_shift_norm=use_scale_shift_norm, ), AttentionBlock( ch, use_checkpoint=use_checkpoint, num_heads=num_heads, num_head_channels=num_head_channels, use_new_attention_order=use_new_attention_order, ), ResBlock( ch, time_embed_dim, dropout, dims=dims, use_checkpoint=use_checkpoint, use_scale_shift_norm=use_scale_shift_norm, ), ) self._feature_size += ch self.pool = pool if pool == "adaptive": self.out = nn.Sequential( normalization(ch), nn.SiLU(), nn.AdaptiveAvgPool2d((1, 1)), zero_module(conv_nd(dims, ch, out_channels, 1)), nn.Flatten(), ) elif pool == "attention": assert num_head_channels != -1 self.out = nn.Sequential( normalization(ch), nn.SiLU(), AttentionPool2d( (image_size // ds), ch, num_head_channels, out_channels ), ) elif pool == "spatial": self.out = nn.Sequential( nn.Linear(self._feature_size, 2048), nn.ReLU(), nn.Linear(2048, self.out_channels), ) elif pool == "spatial_v2": self.out = nn.Sequential( nn.Linear(self._feature_size, 2048), normalization(2048), nn.SiLU(), nn.Linear(2048, self.out_channels), ) else: raise NotImplementedError(f"Unexpected {pool} pooling") def convert_to_fp16(self): """ Convert the torso of the model to float16. """ self.input_blocks.apply(convert_module_to_f16) self.middle_block.apply(convert_module_to_f16) def convert_to_fp32(self): """ Convert the torso of the model to float32. """ self.input_blocks.apply(convert_module_to_f32) self.middle_block.apply(convert_module_to_f32) def forward(self, x, timesteps): """ Apply the model to an input batch. :param x: an [N x C x ...] Tensor of inputs. :param timesteps: a 1-D batch of timesteps. :return: an [N x K] Tensor of outputs. """ emb = self.time_embed(timestep_embedding(timesteps, self.model_channels)) results = [] h = x.type(self.dtype) for module in self.input_blocks: h = module(h, emb) if self.pool.startswith("spatial"): results.append(h.type(x.dtype).mean(dim=(2, 3))) h = self.middle_block(h, emb) if self.pool.startswith("spatial"): results.append(h.type(x.dtype).mean(dim=(2, 3))) h = th.cat(results, axis=-1) return self.out(h) else: h = h.type(x.dtype) return self.out(h) ================================================ FILE: main.py ================================================ import torch from PIL import Image from collections import defaultdict import torch.nn.functional as TF import torchvision.datasets as dsets from torchvision import transforms import numpy as np import torch.optim as optim import torchvision import matplotlib.pyplot as plt import utils from guided_diffusion.unet import UNetModel import math from tensorboardX import SummaryWriter import os import json from collections import namedtuple import argparse from torchvision.utils import save_image from tqdm import tqdm from blur_diffusion import Deblurring, ForwardBlurIncreasing, gaussian_kernel_1d from utils import normalize_np, clear from EMA import EMA from torch.nn import DataParallel from fid import FID from scipy.integrate import solve_ivp parser = argparse.ArgumentParser(description='Configs') parser.add_argument('--gpu', type=str, help='gpu num') parser.add_argument('--dataset', type=str, help='cifar10 / mnist') parser.add_argument('--name', type=str, help='Saving directory name') parser.add_argument('--ckpt', default='', type=str, help='UNet checkpoint') parser.add_argument('--bsize', default=16, type=int, help='batchsize') parser.add_argument('--N', default=1000, type=int, help='Max diffusion timesteps') parser.add_argument('--sig', default=0.4, type=float, help='sigma value for blur kernel') parser.add_argument('--sig_min', default=0, type=float, help='sigma value for blur kernel') parser.add_argument('--sig_max', default=0.1, type=float, help='sigma value for blur kernel') parser.add_argument('--lr', default=0.00005, type=float, help='learning rate') parser.add_argument('--noise_schedule', default='linear', type=str, help='Type of noise schedule to use') parser.add_argument('--betamin', default=0.0001, type=float, help='beta (min). get_score(1) can diverge if this is too low.') parser.add_argument('--betamax', default=0.02, type=float, help='beta (max)') parser.add_argument('--fromprior', default=True, type=bool, help='start sampling from prior') parser.add_argument('--gtscore', action='store_true', help='Use ground truth score for reverse diffusion') parser.add_argument('--max_iter', default=1500000, type=int, help='max iterations') parser.add_argument('--eval_iter', default=10000, type=int, help='eval iterations') parser.add_argument('--fid_iter', default=50000, type=int, help='eval iterations') parser.add_argument('--fid_num_samples', default=10000, type=int, help='eval iterations') parser.add_argument('--fid_bsize', default=32, type=int, help='eval iterations') parser.add_argument('--loss_type', type=str, default = 'eps_simple', choices=['sm_simple', 'eps_simple', 'sm_exact', 'std_matching']) parser.add_argument('--f_type', type=str, default = 'linear', choices=['linear', 'log', 'quadratic', 'cubic', 'quartic', 'triangular']) parser.add_argument('--dropout', default=0, type=float, help='dropout') # EMA, save parser.add_argument('--use_ema', action='store_true', help='use EMA or not') parser.add_argument('--inference', action='store_true') parser.add_argument('--freq_feat', action='store_true', help = "concat Utx_i") parser.add_argument('--ode', action='store_true', help = "ODE fast sampler") parser.add_argument('--ema_decay', type=float, default=0.9999, help='decay rate for EMA') parser.add_argument('--save_every', type=int, default=50000, help='How often we wish to save ckpts') opt = parser.parse_args() dataset = opt.dataset device = torch.device(f'cuda:{opt.gpu}') device = torch.device('cuda') print("N:", opt.N) N = opt.N bsize = opt.bsize beta_min = opt.betamin beta_max = opt.betamax sig = opt.sig if opt.gtscore: opt.fromprior = True if dataset == 'mnist': train_dir = '/home/ubuntu/code/sangyoon/datasets/mnist_train' dataset_train = dsets.MNIST(root='/home/ubuntu/code/sangyoon/datasets/mnist_train', train=True, transform=transforms.Compose([transforms.Resize((32, 32)), transforms.ToTensor()]), download=True) dataset_test = dsets.MNIST(root='/home/ubuntu/code/sangyoon/datasets/mnist_test', train=False, transform=transforms.Compose([transforms.Resize((32, 32)), transforms.ToTensor()]), download=True) elif dataset == 'cifar10': train_dir = "/home/ubuntu/code/sangyoon/datasets/cifar10_train/png" dataset_train = dsets.CIFAR10(root='../datasets/cifar10_train', train=True, transform=transforms.Compose([transforms.RandomHorizontalFlip(0.5), transforms.ToTensor()]), download=True ) dataset_test = dsets.CIFAR10(root='../datasets/cifar10_test', train=False, transform=transforms.Compose([transforms.RandomHorizontalFlip(0.5), transforms.ToTensor()]), download=True ) elif dataset == 'lsun-bedroom': res = 64 train_dir = "/home/ubuntu/code/sangyoon/forward-blur/datasets/bedroom_train/bedroom_train_lmdb/imgs" dataset_train = dsets.LSUN(root='../forward-blur/datasets/bedroom_train', classes=['bedroom_train'], transform=transforms.Compose([transforms.Resize((res,res)), transforms.ToTensor()]) ) dataset_test = dsets.LSUN(root='../forward-blur/datasets/bedroom_train', classes=['bedroom_train'], transform=transforms.Compose([transforms.Resize((res,res)), transforms.ToTensor()]) ) elif dataset == 'lsun-church': res = 64 train_dir = "/home/ubuntu/code/sangyoon/datasets/church_train/church_outdoor_train_lmdb/imgs" dataset_train = dsets.LSUN(root='/home/ubuntu/code/sangyoon/datasets/church_train', classes=['church_outdoor_train'], transform=transforms.Compose([transforms.Resize((res,res)), transforms.ToTensor()]) ) dataset_test = dsets.LSUN(root='/home/ubuntu/code/sangyoon/datasets/church_test', classes=['church_outdoor_val'], transform=transforms.Compose([transforms.Resize((res,res)), transforms.ToTensor()]) ) fid_eval = FID(real_dir = train_dir, device = device) resolution = dataset_train[0][0].shape[-1] input_nc = dataset_train[0][0].shape[0] ksize = resolution * 2 - 1 pad = 0 # define forward blur kernel = gaussian_kernel_1d(ksize, sig) blur = Deblurring(kernel, input_nc, resolution, device=device) print("blur.U_small.shape:", blur.U_small.shape) D_diag = blur.singulars() fb = ForwardBlurIncreasing(N=N, beta_min=beta_min, beta_max=beta_max, sig=sig, sig_max = opt.sig_max, sig_min = opt.sig_min, D_diag=D_diag, blur=blur, channel=input_nc, device=device, noise_schedule=opt.noise_schedule, resolution=resolution, pad=pad, f_type=opt.f_type) dir = os.path.join('experiments', opt.name) writer = SummaryWriter(dir) # # Bedroom config - ADM # model_channels = 256 # num_res_blocks = 2 # channel_mult = (1, 1, 2, 2, 4, 4) # attention_resolutions = [8,16,32] # out_channels = 6 # mean and varaince # dropout = 0.1 # resblock_updown = True # # bsize 256 # # iter -- cat (200K), horse (250K), bedroom (500K) # lr = 1e-5 # Fine tuning # model = UNetModel(image_size = resolution, in_channels = input_nc, model_channels = model_channels, out_channels = out_channels, # num_res_blocks=num_res_blocks, channel_mult=channel_mult, attention_resolutions=attention_resolutions, # dropout = dropout, resblock_updown = resblock_updown) model = UNetModel(resolution, input_nc, 128, input_nc, blur = blur, dropout=opt.dropout, freq_feat = opt.freq_feat) if not opt.ckpt == '' and os.path.exists(opt.ckpt): model.load_state_dict(torch.load(opt.ckpt)) if torch.cuda.device_count() > 1: print("Let's use", torch.cuda.device_count(), "GPUs!") # dim = 0 [30, xxx] -> [10, ...], [10, ...], [10, ...] on 3 GPUs model = DataParallel(model) model.to(device) print("input_nc", input_nc, "resolution", resolution) data_loader = torch.utils.data.DataLoader(dataset=dataset_train, batch_size=bsize, shuffle=True, drop_last=True) data_loader_test = torch.utils.data.DataLoader(dataset=dataset_test, batch_size=bsize, shuffle=False, drop_last=True) optimizer = optim.Adam(model.parameters(), lr=opt.lr) if opt.use_ema: optimizer = EMA(optimizer, ema_decay=opt.ema_decay) # forward process visualization sample = dataset_train[1][0].unsqueeze(0) x_0 = sample[:4] x_0 = x_0.to(device) i = np.array([500] * x_0.shape[0]) i = torch.from_numpy(i).to(device) fb.sanity(x_0, i) sample_list = [] for i in range(0, N+1, N//10): if i == 0: sample_list.append(x_0[0]) continue i = np.array([i] * x_0.shape[0]) i = torch.from_numpy(i).to(device) x_i = fb.get_x_i(x_0, i) sample_list.append(x_i[0]) print(f"x_{i.item()}.std() = {x_i.std()}") print(f"x_{i.item()}.mean() = {x_i.mean()}") grid_sample = torch.cat(sample_list, dim=2) utils.tensor_imsave(grid_sample, "./" + dir, "forward_process.jpg") with open(os.path.join(dir, "config.json"), "w") as json_file: json.dump(vars(opt), json_file) import time meta_iter = 0 for step in range(opt.max_iter): if not opt.inference: elips = time.time() try: x_0, _ = train_iter.next() except: train_iter = iter(data_loader) image, _ = train_iter.next() """ training """ assert x_0.shape[-1] == resolution, f"{x_0.shape}" i = np.random.uniform(1 / N, 1, size = (x_0.shape[0])) * N i = torch.from_numpy(i).to(device).type(torch.long) x_0 = x_0.to(device) x_i, eps = fb.get_x_i(x_0, i, return_eps = True) if opt.loss_type == "sm_simple": loss = fb.get_loss_i_simple(model, x_0, x_i, i) elif opt.loss_type == "eps_simple": loss = fb.get_loss_i_eps_simple(model, x_i, i, eps) elif opt.loss_type == "sm_exact": loss = fb.get_loss_i_exact(model, x_0, x_i, i) elif opt.loss_type == "std_matching": loss = fb.get_loss_i_std_matching(model, x_i, i, eps) writer.add_scalar('loss_train', loss, step) loss.backward() optimizer.step() optimizer.zero_grad() if step % 100 == 0: print(step, loss) # print(f"time: {time.time() - elips}") # Calcuate FID if step > 240001: fid_iter = opt.fid_iter else: fid_iter = 240000 if (step % fid_iter == 0 and step > 0): id = 0 if not os.path.exists(os.path.join("./",dir, f"{step}")): os.mkdir(os.path.join("./",dir, f"{step}")) with torch.no_grad(): if opt.use_ema: optimizer.swap_parameters_with_ema(store_params_in_ema=True) model = model.eval() for _ in range(opt.fid_num_samples // opt.fid_bsize): i = np.array([opt.N - 1] * opt.fid_bsize) i = torch.from_numpy(i).to(device) pred = fb.get_x_N([opt.fid_bsize, input_nc, resolution, resolution], i) for i in reversed(range(1, opt.N)): i = np.array([i] * opt.fid_bsize) i = torch.from_numpy(i).to(device) if opt.loss_type == "sm_simple": s = model(pred, i) elif opt.loss_type == "eps_simple": eps = model(pred, i) s = fb.get_score_from_eps(eps, i) elif opt.loss_type == "sm_exact": s = model(pred, i) elif opt.loss_type == "std_matching": std = model(pred, i) s = fb.get_score_from_std(std, i) else: raise NotImplementedError s = fb.U_I_minus_B_Ut(s, i) rms = lambda x: torch.sqrt(torch.mean(x ** 2)) # print(f"rms(s) * fb._beta_i(i) = {rms(s) * fb._beta_i(i)[0]}") hf = pred - fb.W(pred, i) # Anderson theorem pred1 = pred + hf # unsharpening mask filtering pred2 = pred1 + s # # denoising if i[0] > 2: pred = pred2 + fb.U_I_minus_B_sqrt_Ut(torch.randn_like(pred), i) # inject noise else: pred = pred2 # print(f"i = {i[0]}, rmse = {torch.sqrt(torch.mean(pred**2))}, mean = {torch.mean(pred)} std = {torch.std(pred)}" ) for sample in pred: save_image(sample, os.path.join(dir, f"{step}", f"{id:05d}.png")) id += 1 if opt.use_ema: optimizer.swap_parameters_with_ema(store_params_in_ema=True) model = model.train() fid = fid_eval(os.path.join(dir, f"{step}")) writer.add_scalar('fid', fid, step) print(f"step {step}, fid = {fid}") if (step % opt.eval_iter == 0 and step > 0) or opt.inference: """ sampling (eval) """ cnt = 0 loss = 0 with torch.no_grad(): if opt.use_ema: optimizer.swap_parameters_with_ema(store_params_in_ema=True) model = model.eval() if opt.ode: raise NotImplementedError def to_flattened_numpy(x): """Flatten a torch tensor `x` and convert it to numpy.""" return x.detach().cpu().numpy().reshape((-1,)) def from_flattened_numpy(x, shape): """Form a torch tensor with the given `shape` from a flattened numpy array `x`.""" return torch.from_numpy(x.reshape(shape)) def ode_func(i, y): i = int(i*N) print(f"i = {i}") y = from_flattened_numpy(y, [bsize, input_nc, resolution, resolution]).to(device).type(torch.float32) i = np.array([N - 1] * bsize) i = torch.from_numpy(i).to(device) if opt.loss_type == "sm_simple": s = model(y, i) elif opt.loss_type == "eps_simple": eps = model(y, i) s = fb.get_score_from_eps(eps, i) elif opt.loss_type == "sm_exact": s = model(y, i) elif opt.loss_type == "std_matching": std = model(y, i) s = fb.get_score_from_std(std, i) else: raise NotImplementedError s = fb.U_I_minus_B_Ut(s, i) hf = y - fb.W(y, i) dt = - 1.0 / N drift = (s/2 + hf) / dt drift = to_flattened_numpy(drift) return drift x_N = fb.get_x_N([bsize, input_nc, resolution, resolution], N) solution = solve_ivp(ode_func, (1, 1e-3), to_flattened_numpy(x_N), rtol=1e-3, atol=1e-3, method="RK45") nfe = solution.nfev solution = torch.tensor(solution.y[:, -1]).reshape(x_N.shape).to(device).type(torch.float32) save_image(solution, "./solution.jpg") print(f"nfe = {nfe}") raise NotImplementedError for x_0, _ in data_loader_test: x_0 = x_0.to(device) # for v in range(0, 250, 20): # x_0[:, :, v, :] = 0 if opt.fromprior: i = np.array([N - 1] * x_0.shape[0]) i = torch.from_numpy(i).to(device) pred = fb.get_x_N(x_0.shape, i) print(f"pred.std() = {pred.std()}") else: i = np.array([N-1] * x_0.shape[0]) i = torch.from_numpy(i).to(device) pred = fb.get_x_i(x_0, i) preds = [pred] for i in reversed(range(1, N)): i = np.array([i] * x_0.shape[0]) i = torch.from_numpy(i).to(device) if opt.gtscore: s = fb.get_score_gt(pred, x_0, i) else: if opt.loss_type == "sm_simple": s = model(pred, i) elif opt.loss_type == "eps_simple": eps = model(pred, i) s = fb.get_score_from_eps(eps, i) elif opt.loss_type == "sm_exact": s = model(pred, i) elif opt.loss_type == "std_matching": std = model(pred, i) s = fb.get_score_from_std(std, i) else: raise NotImplementedError s = fb.U_I_minus_B_Ut(s, i) rms = lambda x: torch.sqrt(torch.mean(x ** 2)) # print(f"rms(s) * fb._beta_i(i) = {rms(s) * fb._beta_i(i)[0]}") hf = pred - fb.W(pred, i) # Anderson theorem pred1 = pred + hf # unsharpening mask filtering pred2 = pred1 + s # # denoising if i[0] > 2: pred = pred2 + fb.U_I_minus_B_sqrt_Ut(torch.randn_like(pred), i) # inject noise else: pred = pred2 # print(f"i = {i[0]}, rmse = {torch.sqrt(torch.mean(pred**2))}, mean = {torch.mean(pred)} std = {torch.std(pred)}") # assert rms(pred) < 100 if (i[0]) % (N // 10) == 0: img = pred[0] preds.append(pred) preds.append(pred) assert x_0.shape == pred.shape # visualize grid = torch.cat(preds, dim=3) # grid_sample.shape: (bsize, channel, H, W * 12) # (batch_size, channel, H, W * 12) -> (channel, H * bsize, W * 12) grid = grid.permute(1, 0, 2, 3).contiguous().view(grid.shape[1], -1, grid.shape[3]) # (bsize, channel, H, W) -> (channel, H, W * bsize) gt = x_0.permute(1, 2, 0, 3).contiguous().view(x_0.shape[1], -1, x_0.shape[3] * x_0.shape[0]) if cnt <= 2: utils.tensor_imsave(gt, "./" + dir, f"{step}_{cnt}_GT.jpg") utils.tensor_imsave(grid, "./" + dir, f"{step}_{cnt}_pred.jpg") cnt += 1 loss += TF.l1_loss(x_0, pred) / 2 if cnt == 2: break print(f"step: {step} loss: {loss}") writer.add_scalar('loss_val', loss, meta_iter) f = open('./' + str(dir) + '/log.txt', 'a') f.write(f"Step: {step} loss: {loss}" + '\n') f.close() model.train() if opt.use_ema: optimizer.swap_parameters_with_ema(store_params_in_ema=True) if step % opt.save_every == 1: if opt.use_ema: optimizer.swap_parameters_with_ema(store_params_in_ema=True) if torch.cuda.device_count() > 1: torch.save(model.module.state_dict(), os.path.join(dir, f"model_{step}.ckpt")) else: torch.save(model.state_dict(), os.path.join(dir, f"model_{step}.ckpt")) if opt.use_ema: optimizer.swap_parameters_with_ema(store_params_in_ema=True) ================================================ FILE: pytorch_fid/__init__.py ================================================ """ https://github.com/mseitzer/pytorch-fid/blob/master/src/pytorch_fid/fid_score.py """ __version__ = '0.2.1' ================================================ FILE: pytorch_fid/__main__.py ================================================ import pytorch_fid.fid_score pytorch_fid.fid_score.main() ================================================ FILE: pytorch_fid/fid_score.py ================================================ """Calculates the Frechet Inception Distance (FID) to evalulate GANs The FID metric calculates the distance between two distributions of images. Typically, we have summary statistics (mean & covariance matrix) of one of these distributions, while the 2nd distribution is given by a GAN. When run as a stand-alone program, it compares the distribution of images that are stored as PNG/JPEG at a specified location with a distribution given by summary statistics (in pickle format). The FID is calculated by assuming that X_1 and X_2 are the activations of the pool_3 layer of the inception net for generated samples and real world samples respectively. See --help to see further details. Code apapted from https://github.com/bioinf-jku/TTUR to use PyTorch instead of Tensorflow Copyright 2018 Institute of Bioinformatics, JKU Linz Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. """ import os import pathlib from argparse import ArgumentDefaultsHelpFormatter, ArgumentParser import numpy as np import torch import torchvision.transforms as TF from PIL import Image from scipy import linalg from torch.nn.functional import adaptive_avg_pool2d try: from tqdm import tqdm except ImportError: # If tqdm is not available, provide a mock version of it def tqdm(x): return x try: from pytorch_fid.inception import InceptionV3 except: from inception import InceptionV3 parser = ArgumentParser(formatter_class=ArgumentDefaultsHelpFormatter) parser.add_argument('--batch-size', type=int, default=50, help='Batch size to use') parser.add_argument('--num-workers', type=int, help=('Number of processes to use for data loading. ' 'Defaults to `min(8, num_cpus)`')) parser.add_argument('--device', type=str, default=None, help='Device to use. Like cuda, cuda:0 or cpu') parser.add_argument('--dims', type=int, default=2048, choices=list(InceptionV3.BLOCK_INDEX_BY_DIM), help=('Dimensionality of Inception features to use. ' 'By default, uses pool3 features')) parser.add_argument('path', type=str, nargs=2, help=('Paths to the generated images or ' 'to .npz statistic files')) IMAGE_EXTENSIONS = {'bmp', 'jpg', 'jpeg', 'pgm', 'png', 'ppm', 'tif', 'tiff', 'webp'} class ImagePathDataset(torch.utils.data.Dataset): def __init__(self, files, transforms=None): self.files = files self.transforms = transforms def __len__(self): return len(self.files) def __getitem__(self, i): path = self.files[i] img = Image.open(path).convert('RGB') if self.transforms is not None: img = self.transforms(img) return img def get_activations(files, model, batch_size=50, dims=2048, device='cpu', num_workers=1): """Calculates the activations of the pool_3 layer for all images. Params: -- files : List of image files paths -- model : Instance of inception model -- batch_size : Batch size of images for the model to process at once. Make sure that the number of samples is a multiple of the batch size, otherwise some samples are ignored. This behavior is retained to match the original FID score implementation. -- dims : Dimensionality of features returned by Inception -- device : Device to run calculations -- num_workers : Number of parallel dataloader workers Returns: -- A numpy array of dimension (num images, dims) that contains the activations of the given tensor when feeding inception with the query tensor. """ model.eval() if batch_size > len(files): print(('Warning: batch size is bigger than the data size. ' 'Setting batch size to data size')) batch_size = len(files) dataset = ImagePathDataset(files, transforms=TF.ToTensor()) dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=False, drop_last=False, num_workers=num_workers) pred_arr = np.empty((len(files), dims)) start_idx = 0 for batch in tqdm(dataloader): batch = batch.to(device) with torch.no_grad(): pred = model(batch)[0] # If model output is not scalar, apply global spatial average pooling. # This happens if you choose a dimensionality not equal 2048. if pred.size(2) != 1 or pred.size(3) != 1: pred = adaptive_avg_pool2d(pred, output_size=(1, 1)) pred = pred.squeeze(3).squeeze(2).cpu().numpy() pred_arr[start_idx:start_idx + pred.shape[0]] = pred start_idx = start_idx + pred.shape[0] return pred_arr def calculate_frechet_distance(mu1, sigma1, mu2, sigma2, eps=1e-6): """Numpy implementation of the Frechet Distance. The Frechet distance between two multivariate Gaussians X_1 ~ N(mu_1, C_1) and X_2 ~ N(mu_2, C_2) is d^2 = ||mu_1 - mu_2||^2 + Tr(C_1 + C_2 - 2*sqrt(C_1*C_2)). Stable version by Dougal J. Sutherland. Params: -- mu1 : Numpy array containing the activations of a layer of the inception net (like returned by the function 'get_predictions') for generated samples. -- mu2 : The sample mean over activations, precalculated on an representative data set. -- sigma1: The covariance matrix over activations for generated samples. -- sigma2: The covariance matrix over activations, precalculated on an representative data set. Returns: -- : The Frechet Distance. """ mu1 = np.atleast_1d(mu1) mu2 = np.atleast_1d(mu2) sigma1 = np.atleast_2d(sigma1) sigma2 = np.atleast_2d(sigma2) assert mu1.shape == mu2.shape, \ 'Training and test mean vectors have different lengths' assert sigma1.shape == sigma2.shape, \ 'Training and test covariances have different dimensions' diff = mu1 - mu2 # Product might be almost singular covmean, _ = linalg.sqrtm(sigma1.dot(sigma2), disp=False) if not np.isfinite(covmean).all(): msg = ('fid calculation produces singular product; ' 'adding %s to diagonal of cov estimates') % eps print(msg) offset = np.eye(sigma1.shape[0]) * eps covmean = linalg.sqrtm((sigma1 + offset).dot(sigma2 + offset)) # Numerical error might give slight imaginary component if np.iscomplexobj(covmean): if not np.allclose(np.diagonal(covmean).imag, 0, atol=1e-3): m = np.max(np.abs(covmean.imag)) raise ValueError('Imaginary component {}'.format(m)) covmean = covmean.real tr_covmean = np.trace(covmean) return (diff.dot(diff) + np.trace(sigma1) + np.trace(sigma2) - 2 * tr_covmean) def calculate_activation_statistics(files, model, batch_size=50, dims=2048, device='cpu', num_workers=1): """Calculation of the statistics used by the FID. Params: -- files : List of image files paths -- model : Instance of inception model -- batch_size : The images numpy array is split into batches with batch size batch_size. A reasonable batch size depends on the hardware. -- dims : Dimensionality of features returned by Inception -- device : Device to run calculations -- num_workers : Number of parallel dataloader workers Returns: -- mu : The mean over samples of the activations of the pool_3 layer of the inception model. -- sigma : The covariance matrix of the activations of the pool_3 layer of the inception model. """ act = get_activations(files, model, batch_size, dims, device, num_workers) mu = np.mean(act, axis=0) sigma = np.cov(act, rowvar=False) return mu, sigma def compute_statistics_of_path(path, model, batch_size, dims, device, num_workers=1): if path.endswith('.npz'): with np.load(path) as f: m, s = f['mu'][:], f['sigma'][:] else: path = pathlib.Path(path) files = sorted([file for ext in IMAGE_EXTENSIONS for file in path.glob('*.{}'.format(ext))]) print(f"path: {path}") m, s = calculate_activation_statistics(files, model, batch_size, dims, device, num_workers) # Save to a npz file np.savez_compressed(path/'inception_statistics.npz', mu=m, sigma=s) return m, s def calculate_fid_given_paths(paths, batch_size, device, dims, num_workers=1): """Calculates the FID of two paths""" for p in paths: if not os.path.exists(p): raise RuntimeError('Invalid path: %s' % p) block_idx = InceptionV3.BLOCK_INDEX_BY_DIM[dims] model = InceptionV3([block_idx]).to(device) m1, s1 = compute_statistics_of_path(paths[0], model, batch_size, dims, device, num_workers) m2, s2 = compute_statistics_of_path(paths[1], model, batch_size, dims, device, num_workers) fid_value = calculate_frechet_distance(m1, s1, m2, s2) return fid_value def main(): args = parser.parse_args() if args.device is None: device = torch.device('cuda' if (torch.cuda.is_available()) else 'cpu') else: device = torch.device(args.device) if args.num_workers is None: num_avail_cpus = len(os.sched_getaffinity(0)) num_workers = min(num_avail_cpus, 8) else: num_workers = args.num_workers fid_value = calculate_fid_given_paths(args.path, args.batch_size, device, args.dims, num_workers) print('FID: ', fid_value) if __name__ == '__main__': main() ================================================ FILE: pytorch_fid/inception.py ================================================ import torch import torch.nn as nn import torch.nn.functional as F import torchvision try: from torchvision.models.utils import load_state_dict_from_url except ImportError: from torch.utils.model_zoo import load_url as load_state_dict_from_url # Inception weights ported to Pytorch from # http://download.tensorflow.org/models/image/imagenet/inception-2015-12-05.tgz FID_WEIGHTS_URL = 'https://github.com/mseitzer/pytorch-fid/releases/download/fid_weights/pt_inception-2015-12-05-6726825d.pth' # noqa: E501 class InceptionV3(nn.Module): """Pretrained InceptionV3 network returning feature maps""" # Index of default block of inception to return, # corresponds to output of final average pooling DEFAULT_BLOCK_INDEX = 3 # Maps feature dimensionality to their output blocks indices BLOCK_INDEX_BY_DIM = { 64: 0, # First max pooling features 192: 1, # Second max pooling featurs 768: 2, # Pre-aux classifier features 2048: 3 # Final average pooling features } def __init__(self, output_blocks=(DEFAULT_BLOCK_INDEX,), resize_input=True, normalize_input=True, requires_grad=False, use_fid_inception=True): """Build pretrained InceptionV3 Parameters ---------- output_blocks : list of int Indices of blocks to return features of. Possible values are: - 0: corresponds to output of first max pooling - 1: corresponds to output of second max pooling - 2: corresponds to output which is fed to aux classifier - 3: corresponds to output of final average pooling resize_input : bool If true, bilinearly resizes input to width and height 299 before feeding input to model. As the network without fully connected layers is fully convolutional, it should be able to handle inputs of arbitrary size, so resizing might not be strictly needed normalize_input : bool If true, scales the input from range (0, 1) to the range the pretrained Inception network expects, namely (-1, 1) requires_grad : bool If true, parameters of the model require gradients. Possibly useful for finetuning the network use_fid_inception : bool If true, uses the pretrained Inception model used in Tensorflow's FID implementation. If false, uses the pretrained Inception model available in torchvision. The FID Inception model has different weights and a slightly different structure from torchvision's Inception model. If you want to compute FID scores, you are strongly advised to set this parameter to true to get comparable results. """ super(InceptionV3, self).__init__() self.resize_input = resize_input self.normalize_input = normalize_input self.output_blocks = sorted(output_blocks) self.last_needed_block = max(output_blocks) assert self.last_needed_block <= 3, \ 'Last possible output block index is 3' self.blocks = nn.ModuleList() if use_fid_inception: inception = fid_inception_v3() else: inception = _inception_v3(pretrained=True) # Block 0: input to maxpool1 block0 = [ inception.Conv2d_1a_3x3, inception.Conv2d_2a_3x3, inception.Conv2d_2b_3x3, nn.MaxPool2d(kernel_size=3, stride=2) ] self.blocks.append(nn.Sequential(*block0)) # Block 1: maxpool1 to maxpool2 if self.last_needed_block >= 1: block1 = [ inception.Conv2d_3b_1x1, inception.Conv2d_4a_3x3, nn.MaxPool2d(kernel_size=3, stride=2) ] self.blocks.append(nn.Sequential(*block1)) # Block 2: maxpool2 to aux classifier if self.last_needed_block >= 2: block2 = [ inception.Mixed_5b, inception.Mixed_5c, inception.Mixed_5d, inception.Mixed_6a, inception.Mixed_6b, inception.Mixed_6c, inception.Mixed_6d, inception.Mixed_6e, ] self.blocks.append(nn.Sequential(*block2)) # Block 3: aux classifier to final avgpool if self.last_needed_block >= 3: block3 = [ inception.Mixed_7a, inception.Mixed_7b, inception.Mixed_7c, nn.AdaptiveAvgPool2d(output_size=(1, 1)) ] self.blocks.append(nn.Sequential(*block3)) for param in self.parameters(): param.requires_grad = requires_grad def forward(self, inp): """Get Inception feature maps Parameters ---------- inp : torch.autograd.Variable Input tensor of shape Bx3xHxW. Values are expected to be in range (0, 1) Returns ------- List of torch.autograd.Variable, corresponding to the selected output block, sorted ascending by index """ outp = [] x = inp if self.resize_input: x = F.interpolate(x, size=(299, 299), mode='bilinear', align_corners=False) if self.normalize_input: x = 2 * x - 1 # Scale from range (0, 1) to range (-1, 1) for idx, block in enumerate(self.blocks): x = block(x) if idx in self.output_blocks: outp.append(x) if idx == self.last_needed_block: break return outp def _inception_v3(*args, **kwargs): """Wraps `torchvision.models.inception_v3` Skips default weight inititialization if supported by torchvision version. See https://github.com/mseitzer/pytorch-fid/issues/28. """ try: version = tuple(map(int, torchvision.__version__.split('.')[:2])) except ValueError: # Just a caution against weird version strings version = (0,) if version >= (0, 6): kwargs['init_weights'] = False return torchvision.models.inception_v3(*args, **kwargs) def fid_inception_v3(): """Build pretrained Inception model for FID computation The Inception model for FID computation uses a different set of weights and has a slightly different structure than torchvision's Inception. This method first constructs torchvision's Inception and then patches the necessary parts that are different in the FID Inception model. """ inception = _inception_v3(num_classes=1008, aux_logits=False, pretrained=False) inception.Mixed_5b = FIDInceptionA(192, pool_features=32) inception.Mixed_5c = FIDInceptionA(256, pool_features=64) inception.Mixed_5d = FIDInceptionA(288, pool_features=64) inception.Mixed_6b = FIDInceptionC(768, channels_7x7=128) inception.Mixed_6c = FIDInceptionC(768, channels_7x7=160) inception.Mixed_6d = FIDInceptionC(768, channels_7x7=160) inception.Mixed_6e = FIDInceptionC(768, channels_7x7=192) inception.Mixed_7b = FIDInceptionE_1(1280) inception.Mixed_7c = FIDInceptionE_2(2048) state_dict = load_state_dict_from_url(FID_WEIGHTS_URL, progress=True) inception.load_state_dict(state_dict) return inception class FIDInceptionA(torchvision.models.inception.InceptionA): """InceptionA block patched for FID computation""" def __init__(self, in_channels, pool_features): super(FIDInceptionA, self).__init__(in_channels, pool_features) def forward(self, x): branch1x1 = self.branch1x1(x) branch5x5 = self.branch5x5_1(x) branch5x5 = self.branch5x5_2(branch5x5) branch3x3dbl = self.branch3x3dbl_1(x) branch3x3dbl = self.branch3x3dbl_2(branch3x3dbl) branch3x3dbl = self.branch3x3dbl_3(branch3x3dbl) # Patch: Tensorflow's average pool does not use the padded zero's in # its average calculation branch_pool = F.avg_pool2d(x, kernel_size=3, stride=1, padding=1, count_include_pad=False) branch_pool = self.branch_pool(branch_pool) outputs = [branch1x1, branch5x5, branch3x3dbl, branch_pool] return torch.cat(outputs, 1) class FIDInceptionC(torchvision.models.inception.InceptionC): """InceptionC block patched for FID computation""" def __init__(self, in_channels, channels_7x7): super(FIDInceptionC, self).__init__(in_channels, channels_7x7) def forward(self, x): branch1x1 = self.branch1x1(x) branch7x7 = self.branch7x7_1(x) branch7x7 = self.branch7x7_2(branch7x7) branch7x7 = self.branch7x7_3(branch7x7) branch7x7dbl = self.branch7x7dbl_1(x) branch7x7dbl = self.branch7x7dbl_2(branch7x7dbl) branch7x7dbl = self.branch7x7dbl_3(branch7x7dbl) branch7x7dbl = self.branch7x7dbl_4(branch7x7dbl) branch7x7dbl = self.branch7x7dbl_5(branch7x7dbl) # Patch: Tensorflow's average pool does not use the padded zero's in # its average calculation branch_pool = F.avg_pool2d(x, kernel_size=3, stride=1, padding=1, count_include_pad=False) branch_pool = self.branch_pool(branch_pool) outputs = [branch1x1, branch7x7, branch7x7dbl, branch_pool] return torch.cat(outputs, 1) class FIDInceptionE_1(torchvision.models.inception.InceptionE): """First InceptionE block patched for FID computation""" def __init__(self, in_channels): super(FIDInceptionE_1, self).__init__(in_channels) def forward(self, x): branch1x1 = self.branch1x1(x) branch3x3 = self.branch3x3_1(x) branch3x3 = [ self.branch3x3_2a(branch3x3), self.branch3x3_2b(branch3x3), ] branch3x3 = torch.cat(branch3x3, 1) branch3x3dbl = self.branch3x3dbl_1(x) branch3x3dbl = self.branch3x3dbl_2(branch3x3dbl) branch3x3dbl = [ self.branch3x3dbl_3a(branch3x3dbl), self.branch3x3dbl_3b(branch3x3dbl), ] branch3x3dbl = torch.cat(branch3x3dbl, 1) # Patch: Tensorflow's average pool does not use the padded zero's in # its average calculation branch_pool = F.avg_pool2d(x, kernel_size=3, stride=1, padding=1, count_include_pad=False) branch_pool = self.branch_pool(branch_pool) outputs = [branch1x1, branch3x3, branch3x3dbl, branch_pool] return torch.cat(outputs, 1) class FIDInceptionE_2(torchvision.models.inception.InceptionE): """Second InceptionE block patched for FID computation""" def __init__(self, in_channels): super(FIDInceptionE_2, self).__init__(in_channels) def forward(self, x): branch1x1 = self.branch1x1(x) branch3x3 = self.branch3x3_1(x) branch3x3 = [ self.branch3x3_2a(branch3x3), self.branch3x3_2b(branch3x3), ] branch3x3 = torch.cat(branch3x3, 1) branch3x3dbl = self.branch3x3dbl_1(x) branch3x3dbl = self.branch3x3dbl_2(branch3x3dbl) branch3x3dbl = [ self.branch3x3dbl_3a(branch3x3dbl), self.branch3x3dbl_3b(branch3x3dbl), ] branch3x3dbl = torch.cat(branch3x3dbl, 1) # Patch: The FID Inception model uses max pooling instead of average # pooling. This is likely an error in this specific Inception # implementation, as other Inception models use average pooling here # (which matches the description in the paper). branch_pool = F.max_pool2d(x, kernel_size=3, stride=1, padding=1) branch_pool = self.branch_pool(branch_pool) outputs = [branch1x1, branch3x3, branch3x3dbl, branch_pool] return torch.cat(outputs, 1) ================================================ FILE: train.sh ================================================ export OMP_NUM_THREADS=1 export CUDA_VISIBLE_DEVICES=0 python main.py \ --name test \ --dataset lsun-bedroom \ --use_ema \ --ema_decay 0.9999 \ --noise_schedule linear \ --f_type quartic \ --gpu 0 \ --bsize 16 \ --sig 0.4 \ --sig_min 0 \ --sig_max 0.15 \ --fid_bsize 16 \ --loss_type eps_simple \ --lr 0.00005 \ ================================================ FILE: utils.py ================================================ from torchvision import transforms from PIL import Image import os from torch import nn import torch import glob import numpy as np import torch.nn.functional as F def tensor2pil(t): img = transforms.functional.to_pil_image(denorm(t)) return img def denorm(t): return t*0.5 + 0.5 def tensor_imsave(t, path, fname, denormalization=False, prt=False): # Save tensor as .png file check_folder(path) if denormalization: t = denorm(t) t = torch.clamp(input=t, min=0., max=1.) img = transforms.functional.to_pil_image(t.detach().cpu()) img.save(os.path.join(path,fname)) if prt: print(f"Saved to {os.path.join(path,fname)}") def pil_loader(path): # open path as file to avoid ResourceWarning (https://github.com/python-pillow/Pillow/issues/835) with open(path, 'rb') as f:##Automatic closing using with. rb = read binary img = Image.open(f) return img.convert('RGB') def check_folder(log_dir): if not os.path.exists(log_dir): os.makedirs(log_dir) return log_dir def freeze_model(model): for p in model.parameters(): p.requires_grad = False class TVLoss(nn.Module): def __init__(self,TVLoss_weight=1): super(TVLoss,self).__init__() self.TVLoss_weight = TVLoss_weight def forward(self,x): batch_size = x.size()[0] h_x = x.size()[2] w_x = x.size()[3] count_h = self._tensor_size(x[:,:,1:,:]) count_w = self._tensor_size(x[:,:,:,1:]) h_tv = torch.pow((x[:,:,1:,:]-x[:,:,:h_x-1,:]),2).sum() w_tv = torch.pow((x[:,:,:,1:]-x[:,:,:,:w_x-1]),2).sum() return self.TVLoss_weight*2*(h_tv/count_h+w_tv/count_w)/batch_size def _tensor_size(self,t): return t.size()[1]*t.size()[2]*t.size()[3] def pil_batch_loader(root_path, bsize): # Returns the list of PIL.Images at given path. Length of list is bsize. flist = sorted(glob.glob(root_path, '*.png')) flist = flist[:bsize] result = [] for p in flist: img = pil_loader(p) result.append(img) return result class PSNR(object): def __init__(self, gpu, val_max=1, val_min=0, ycbcr=True): super(PSNR,self).__init__() self.val_max = val_max self.val_min = val_min self.gpu = gpu self.ycbcr = ycbcr def __call__(self,x,y): """ if x.is_cuda: x = x.detach().cpu() if y.is_cuda: y = y.detach().cpu() """ assert len(x.size()) == len(y.size()) with torch.no_grad(): x_lum = rgb_to_ycbcr(x)[:,0] y_lum = rgb_to_ycbcr(y)[:,0] # if len(x.size()) == 3: # mse = torch.mean((y-x)**2) # psnr = 20*torch.log10(torch.tensor(self.val_max-self.val_min, dtype=torch.float).cuda(self.gpu)) - 10*torch.log10(mse) # return psnr # elif len(x.size()) == 4: mse = torch.mean((y_lum-x_lum)**2, dim=[1,2]) psnr = 20*torch.log10(torch.tensor(self.val_max-self.val_min, dtype=torch.float).cuda(self.gpu)) - 10*torch.log10(mse) return torch.mean(psnr) def rgb_to_ycbcr(image): r"""Convert an RGB image to YCbCr. Args: image (torch.Tensor): RGB Image to be converted to YCbCr. Returns: torch.Tensor: YCbCr version of the image. """ if not torch.is_tensor(image): raise TypeError("Input type is not a torch.Tensor. Got {}".format( type(image))) if len(image.shape) < 3 or image.shape[-3] != 3: raise ValueError("Input size must have a shape of (*, 3, H, W). Got {}" .format(image.shape)) r: torch.Tensor = image[..., 0, :, :] g: torch.Tensor = image[..., 1, :, :] b: torch.Tensor = image[..., 2, :, :] delta = .5 y: torch.Tensor = .299 * r + .587 * g + .114 * b cb: torch.Tensor = (b - y) * .564 + delta cr: torch.Tensor = (r - y) * .713 + delta return torch.stack((y, cb, cr), -3) def clear(x): return x.detach().cpu().squeeze().numpy() def normalize_np(x): x -= x.min() x /= x.max() return x def pad(x, pad): return F.pad(x, (pad, pad, pad, pad), 'reflect') def crop(x, pad): if pad == 0: return x else: return x[..., pad:-pad, pad:-pad]