LION/utils/utils.py

1534 lines
58 KiB
Python
Raw Normal View History

2023-01-23 05:14:49 +00:00
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto. Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.
"""copied and modified from https://github.com/NVlabs/LSGM/blob/5eae2f385c014f2250c3130152b6be711f6a3a5a/util/utils.py"""
from loguru import logger
from comet_ml import Experiment, ExistingExperiment
import wandb as WB
import os
import math
import shutil
import json
2023-01-23 05:14:49 +00:00
import time
import sys
import types
from PIL import Image
import torch
import torch.nn as nn
import numpy as np
from torch import optim
import torch.distributed as dist
from torch.cuda.amp import autocast, GradScaler
USE_COMET = int(os.environ.get('USE_COMET', 1))
USE_TFB = int(os.environ.get('USE_TFB', 0))
USE_WB = int(os.environ.get('USE_WB', 0))
print(f'utils/utils.py: USE_COMET={USE_COMET}, USE_WB={USE_WB}')
class PixelNormal(object):
def __init__(self, param, fixed_log_scales=None):
size = param.size()
C = size[1]
if fixed_log_scales is None:
self.num_c = C // 2
# B, 1 or 3, H, W
self.means = param[:, :self.num_c, :, :]
self.log_scales = torch.clamp(
param[:, self.num_c:, :, :], min=-7.0) # B, 1 or 3, H, W
raise NotImplementedError
else:
self.num_c = C
# B, 1 or 3, H, W
self.means = param
# B, 1 or 3, H, W
self.log_scales = view4D(fixed_log_scales, size)
def get_params(self):
return self.means, self.log_scales, self.num_c
def log_prob(self, samples):
B, C, H, W = samples.size()
assert C == self.num_c
log_probs = -0.5 * torch.square(self.means - samples) * torch.exp(-2.0 *
self.log_scales) - self.log_scales - 0.9189385332 # -0.5*log(2*pi)
return log_probs
def sample(self, t=1.):
z, rho = sample_normal_jit(
self.means, torch.exp(self.log_scales)*t) # B, 3, H, W
return z
def log_prob_discrete(self, samples):
"""
Calculates discrete pixel probabilities.
"""
# samples should be in [-1, 1] already
B, C, H, W = samples.size()
assert C == self.num_c
centered = samples - self.means
inv_stdv = torch.exp(- self.log_scales)
plus_in = inv_stdv * (centered + 1. / 255.)
cdf_plus = torch.distributions.Normal(0, 1).cdf(plus_in)
min_in = inv_stdv * (centered - 1. / 255.)
cdf_min = torch.distributions.Normal(0, 1).cdf(min_in)
log_cdf_plus = torch.log(torch.clamp(cdf_plus, min=1e-12))
log_one_minus_cdf_min = torch.log(torch.clamp(1. - cdf_min, min=1e-12))
cdf_delta = cdf_plus - cdf_min
log_probs = torch.where(samples < -0.999, log_cdf_plus, torch.where(samples > 0.999, log_one_minus_cdf_min,
torch.log(torch.clamp(cdf_delta, min=1e-12))))
assert log_probs.size() == samples.size()
return log_probs
def mean(self):
return self.means
class DummyGradScalar(object):
def __init__(self, *args, **kwargs):
pass
def scale(self, input):
return input
def update(self):
pass
def state_dict(self):
return {}
def load_state_dict(self, x):
pass
def step(self, opt):
opt.step()
def unscale_(self, x):
return x
def get_opt(params, cfgopt, use_ema, other_cfg=None):
if cfgopt.type == 'adam':
optimizer = optim.Adam(params,
lr=float(cfgopt.lr),
betas=(cfgopt.beta1, cfgopt.beta2),
weight_decay=cfgopt.weight_decay)
elif cfgopt.type == 'sgd':
optimizer = torch.optim.SGD(params,
lr=float(cfgopt.lr),
momentum=cfgopt.momentum)
elif cfgopt.type == 'adamax':
from utils.adamax import Adamax
logger.info('[Optimizer] Adamax, lr={}, weight_decay={}, eps={}',
cfgopt.lr, cfgopt.weight_decay, 1e-4)
optimizer = Adamax(params, float(cfgopt.lr),
weight_decay=args.weight_decay, eps=1e-4)
else:
assert 0, "Optimizer type should be either 'adam' or 'sgd'"
if use_ema:
logger.info('use_ema')
ema_decay = 0.9999
from .ema import EMA
optimizer = EMA(optimizer, ema_decay=ema_decay)
scheduler = optim.lr_scheduler.LambdaLR(
optimizer, lr_lambda=lambda x: 1.0) # constant lr
scheduler_type = getattr(cfgopt, "scheduler", None)
if scheduler_type is not None and len(scheduler_type) > 0:
logger.info('get scheduler_type: {}', scheduler_type)
if scheduler_type == 'exponential':
decay = float(getattr(cfgopt, "step_decay", 0.1))
scheduler = optim.lr_scheduler.ExponentialLR(optimizer, decay)
elif scheduler_type == 'step':
step_size = int(getattr(cfgopt, "step_epoch", 500))
decay = float(getattr(cfgopt, "step_decay", 0.1))
scheduler = optim.lr_scheduler.StepLR(optimizer,
step_size=step_size,
gamma=decay)
elif scheduler_type == 'linear': # use default setting from shapeLatent
start_epoch = int(getattr(cfgopt, 'sched_start_epoch', 200*1e3))
end_epoch = int(getattr(cfgopt, 'sched_end_epoch', 400*1e3))
end_lr = float(getattr(cfgopt, 'end_lr', 1e-4))
start_lr = cfgopt.lr
def lambda_rule(epoch):
if epoch <= start_epoch:
return 1.0
elif epoch <= end_epoch:
total = end_epoch - start_epoch
delta = epoch - start_epoch
frac = delta / total
return (1 - frac) * 1.0 + frac * (end_lr / start_lr)
else:
return end_lr / start_lr
scheduler = optim.lr_scheduler.LambdaLR(optimizer,
lr_lambda=lambda_rule)
elif scheduler_type == 'lambda': # linear':
step_size = int(getattr(cfgopt, "step_epoch", 2000))
final_ratio = float(getattr(cfgopt, "final_ratio", 0.01))
start_ratio = float(getattr(cfgopt, "start_ratio", 0.5))
duration_ratio = float(getattr(cfgopt, "duration_ratio", 0.45))
def lambda_rule(ep):
lr_l = 1.0 - min(
1,
max(0, ep - start_ratio * step_size) /
float(duration_ratio * step_size)) * (1 - final_ratio)
return lr_l
scheduler = optim.lr_scheduler.LambdaLR(optimizer,
lr_lambda=lambda_rule)
elif scheduler_type == 'cosine_anneal_nocycle':
## logger.info('scheduler_type: {}', scheduler_type)
assert(other_cfg is not None)
final_lr_ratio = float(getattr(cfgopt, "final_lr_ratio", 0.01))
eta_min = float(cfgopt.lr) * final_lr_ratio
eta_max = float(cfgopt.lr)
total_epoch = int(other_cfg.trainer.epochs)
##getattr(cfgopt, "step_epoch", 2000)
start_ratio = float(getattr(cfgopt, "start_ratio", 0.6))
T_max = total_epoch * (1 - start_ratio)
def lambda_rule(ep):
curr_ep = max(0., ep - start_ratio * total_epoch)
lr = eta_min + 0.5 * (eta_max - eta_min) * (
1 + np.cos(np.pi * curr_ep / T_max))
lr_l = lr / eta_max
return lr_l
scheduler = optim.lr_scheduler.LambdaLR(optimizer,
lr_lambda=lambda_rule)
else:
assert 0, "args.schedulers should be either 'exponential' or 'linear' or 'step'"
return optimizer, scheduler
class AvgrageMeter(object):
def __init__(self):
self.reset()
def reset(self):
self.avg = 0
self.sum = 0
self.cnt = 0
def update(self, val, n=1):
self.sum += val * n
self.cnt += n
self.avg = self.sum / self.cnt
class ExpMovingAvgrageMeter(object):
def __init__(self, momentum=0.9):
self.momentum = momentum
self.reset()
def reset(self):
self.avg = 0
def update(self, val):
self.avg = (1. - self.momentum) * self.avg + self.momentum * val
class DummyDDP(nn.Module):
def __init__(self, model):
super(DummyDDP, self).__init__()
self.module = model
def forward(self, *input, **kwargs):
return self.module(*input, **kwargs)
def count_parameters_in_M(model):
return np.sum(np.prod(v.size()) for name, v in model.named_parameters() if "auxiliary" not in name) / 1e6
def save_checkpoint(state, is_best, save):
filename = os.path.join(save, 'checkpoint.pth.tar')
torch.save(state, filename)
if is_best:
best_filename = os.path.join(save, 'model_best.pth.tar')
shutil.copyfile(filename, best_filename)
def save(model, model_path):
torch.save(model.state_dict(), model_path)
def load(model, model_path):
model.load_state_dict(torch.load(model_path))
# def create_exp_dir(path, scripts_to_save=None):
# if not os.path.exists(path):
# os.makedirs(path, exist_ok=True)
# print('Experiment dir : {}'.format(path))
#
# if scripts_to_save is not None:
# if not os.path.exists(os.path.join(path, 'scripts')):
# os.mkdir(os.path.join(path, 'scripts'))
# for script in scripts_to_save:
# dst_file = os.path.join(path, 'scripts', os.path.basename(script))
# shutil.copyfile(script, dst_file)
#
# class Logger(object):
# def __init__(self, rank, save):
# # other libraries may set logging before arriving at this line.
# # by reloading logging, we can get rid of previous configs set by other libraries.
# from importlib import reload
# reload(logging)
# self.rank = rank
# if self.rank == 0:
# log_format = '%(asctime)s %(message)s'
# logging.basicConfig(stream=sys.stdout, level=logging.INFO,
# format=log_format, datefmt='%m/%d %I:%M:%S %p')
# fh = logging.FileHandler(os.path.join(save, 'log.txt'))
# fh.setFormatter(logging.Formatter(log_format))
# logging.getLogger().addHandler(fh)
# self.start_time = time.time()
#
# def info(self, string, *args):
# if self.rank == 0:
# elapsed_time = time.time() - self.start_time
# elapsed_time = time.strftime(
# '(Elapsed: %H:%M:%S) ', time.gmtime(elapsed_time))
# if isinstance(string, str):
# string = elapsed_time + string
# else:
# logging.info(elapsed_time)
# logging.info(string, *args)
def flatten_dict(dd, separator='_', prefix=''):
return {prefix + separator + k if prefix else k: v for kk, vv in dd.items()
for k, v in flatten_dict(vv, separator, kk).items()} \
if isinstance(dd, dict) else {prefix: dd}
class Writer(object):
def __init__(self, rank=0, save=None, exp=None, wandb=False):
self.rank = rank
self.exp = None
self.wandb = False
self.meter_dict = {}
if self.rank == 0:
self.exp = exp
if USE_TFB and save is not None:
logger.info('init TFB: {}', save)
from torch.utils.tensorboard import SummaryWriter
self.writer = SummaryWriter(log_dir=save, flush_secs=20)
else:
logger.info('Not init TFB')
self.writer = None
if self.exp is not None and save is not None:
with open(os.path.join(save, 'url.txt'), 'a') as f:
f.write(self.exp.url)
f.write('\n')
self.wandb = wandb
else:
logger.info('rank={}, init writer as a blackhole', rank)
def set_model_graph(self, *args, **kwargs):
if self.rank == 0 and self.exp is not None:
self.exp.set_model_graph(*args, **kwargs)
@property
def url(self):
if self.exp is not None:
return self.exp.url
else:
return 'none'
def add_hparams(self, cfg, args): # **kwargs):
if self.exp is not None:
self.exp.log_parameters(flatten_dict(cfg))
self.exp.log_parameters(flatten_dict(args))
if self.wandb:
WB.config.update(flatten_dict(cfg))
WB.config.update(flatten_dict(args))
def avg_meter(self, name, value, step=None, epoch=None):
if self.rank == 0:
if name not in self.meter_dict:
self.meter_dict[name] = AvgrageMeter()
self.meter_dict[name].update(value)
def upload_meter(self, step=None, epoch=None):
for name, value in self.meter_dict.items():
self.add_scalar(name, value.avg, step=step, epoch=epoch)
self.meter_dict = {}
def add_scalar(self, *args, **kwargs):
if self.rank == 0 and self.writer is not None:
if 'step' in kwargs:
self.writer.add_scalar(*args,
global_step=kwargs['step'])
else:
self.writer.add_scalar(*args, **kwargs)
if self.exp is not None:
self.exp.log_metric(*args, **kwargs)
if self.wandb:
name = args[0]
v = args[1]
WB.log({name: v})
def log_model(self, name, path):
pass
def log_other(self, name, value):
if self.rank == 0 and self.exp is not None:
self.exp.log_other(name, value)
# if self.rank == 0 and self.exp is not None:
# self.exp.log_model(name, path)
def watch(self, model):
if self.wandb:
WB.watch(model)
def log_points_3d(self, scene_name, points, step=0): # *args, **kwargs):
if self.rank == 0 and self.exp is not None:
self.exp.log_points_3d(*args, **kwargs)
if self.wandb:
WB.log({"point_cloud": WB.Object3D(points)})
def add_figure(self, *args, **kwargs):
if self.rank == 0 and self.writer is not None:
self.writer.add_figure(*args, **kwargs)
def add_image(self, *args, **kwargs):
if self.rank == 0 and self.writer is not None:
self.writer.add_image(*args, **kwargs)
self.writer.flush()
if self.exp is not None:
name, img, i = args
if isinstance(img, Image.Image):
# logger.debug('log PIL Imgae: {}, {}', name, i)
self.exp.log_image(img, name, step=i)
elif type(img) is str:
# logger.debug('log str image: {}, {}: {}', name, i, img)
self.exp.log_image(img, name, step=i)
elif torch.is_tensor(img):
if img.shape[0] in [3, 4] and len(img.shape) == 3: # 3,H,W
img = img.permute(1, 2, 0).contiguous() # 3,H,W -> H,W,3
if img.max() < 100: # [0-1]
ndarr = img.mul(255).add_(0.5).clamp_(
0, 255).to('cpu') # .squeeze()
ndarr = ndarr.numpy().astype(np.uint8)
# .reshape(-1, ndarr.shape[-1]))
im = Image.fromarray(ndarr)
self.exp.log_image(im, name, step=i)
else:
im = img.to('cpu').numpy()
self.exp.log_image(im, name, step=i)
elif isinstance(img, (np.ndarray, np.generic)):
if img.shape[0] == 3 and len(img.shape) == 3: # 3,H,W
img = img.transpose(1, 2, 0)
self.exp.log_image(img, name, step=i)
if self.wandb and torch.is_tensor(img) and self.rank == 0:
## print(img.shape, img.max(), img.type())
WB.log({name: WB.Image(img.numpy())})
def add_histogram(self, *args, **kwargs):
if self.rank == 0 and self.writer is not None:
self.writer.add_histogram(*args, **kwargs)
if self.exp is not None:
name, value, step = args
self.exp.log_histogram_3d(value, name, step)
# *args, **kwargs)
def add_histogram_if(self, write, *args, **kwargs):
if write and False: # Used for debugging.
self.add_histogram(*args, **kwargs)
def close(self, *args, **kwargs):
if self.rank == 0 and self.writer is not None:
self.writer.close()
def log_asset(self, *args, **kwargs):
if self.exp is not None:
self.exp.log_asset(*args, **kwargs)
def common_init(rank, seed, save_dir, comet_key=''):
# we use different seeds per gpu. But we sync the weights after model initialization.
logger.info('[common-init] at rank={}, seed={}', rank, seed)
torch.manual_seed(rank + seed)
np.random.seed(rank + seed)
torch.cuda.manual_seed(rank + seed)
torch.cuda.manual_seed_all(rank + seed)
torch.backends.cudnn.benchmark = True
# prepare logging and tensorboard summary
#logging = Logger(rank, save_dir)
logging = None
if rank == 0:
if os.path.exists('.comet_api'):
comet_args = json.load(open('.comet_api', 'r'))
exp = Experiment(display_summary_level=0,
disabled=USE_COMET == 0,
**comet_args)
exp.set_name(save_dir.split('exp/')[-1])
exp.set_cmd_args()
exp.log_code(folder='./models/')
exp.log_code(folder='./trainers/')
exp.log_code(folder='./utils/')
exp.log_code(folder='./datasets/')
else:
exp = None
if os.path.exists('.wandb_api'):
wb_args = json.load(open('.wandb_api', 'r'))
wb_dir = '../exp/wandb/' if not os.path.exists(
'/workspace/result') else '/workspace/result/wandb/'
if not os.path.exists(wb_dir):
os.makedirs(wb_dir)
WB.init(
project=wb_args['project'],
entity=wb_args['entity'],
name=save_dir.split('exp/')[-1],
dir=wb_dir
)
wandb = True
else:
wandb = False
else:
exp = None
wandb = False
writer = Writer(rank, save_dir, exp, wandb)
logger.info('[common-init] DONE')
return logging, writer
def reduce_tensor(tensor, world_size):
rt = tensor.clone()
dist.all_reduce(rt, op=dist.ReduceOp.SUM)
rt /= world_size
return rt
def get_stride_for_cell_type(cell_type):
if cell_type.startswith('normal') or cell_type.startswith('combiner'):
stride = 1
elif cell_type.startswith('down'):
stride = 2
elif cell_type.startswith('up'):
stride = -1
else:
raise NotImplementedError(cell_type)
return stride
def get_cout(cin, stride):
if stride == 1:
cout = cin
elif stride == -1:
cout = cin // 2
elif stride == 2:
cout = 2 * cin
return cout
def kl_balancer_coeff(num_scales, groups_per_scale, fun='square'):
if fun == 'equal':
coeff = torch.cat([torch.ones(groups_per_scale[num_scales - i - 1])
for i in range(num_scales)], dim=0).cuda()
elif fun == 'linear':
coeff = torch.cat([(2 ** i) * torch.ones(groups_per_scale[num_scales - i - 1]) for i in range(num_scales)],
dim=0).cuda()
elif fun == 'sqrt':
coeff = torch.cat(
[np.sqrt(2 ** i) * torch.ones(groups_per_scale[num_scales - i - 1])
for i in range(num_scales)],
dim=0).cuda()
elif fun == 'square':
coeff = torch.cat(
[np.square(2 ** i) / groups_per_scale[num_scales - i - 1] * torch.ones(groups_per_scale[num_scales - i - 1])
for i in range(num_scales)], dim=0).cuda()
else:
raise NotImplementedError
# convert min to 1.
coeff /= torch.min(coeff)
return coeff
def kl_per_group(kl_all):
kl_vals = torch.mean(kl_all, dim=0)
kl_coeff_i = torch.abs(kl_all)
kl_coeff_i = torch.mean(kl_coeff_i, dim=0, keepdim=True) + 0.01
return kl_coeff_i, kl_vals
def rec_balancer(rec_all, rec_coeff=1.0, npoints=None):
# layer depth increase, alpha_i increase, 1/alpha_i decrease; kl_coeff decrease
# the rec with more points should have higher loss
min_points = min(npoints)
coeff = []
rec_loss = 0
assert(len(rec_all) == len(npoints))
for ni, n in enumerate(npoints):
c = rec_coeff*np.sqrt(n/min_points)
rec_loss += rec_all[ni] * c
coeff.append(c) # the smallest points' loss weight is 1
return rec_loss, coeff, rec_all
def kl_balancer(kl_all, kl_coeff=1.0, kl_balance=False, alpha_i=None):
# layer depth increase, alpha_i increase, 1/alpha_i decrease; kl_coeff decrease
if kl_balance and kl_coeff < 1.0:
alpha_i = alpha_i.unsqueeze(0)
kl_all = torch.stack(kl_all, dim=1)
kl_coeff_i, kl_vals = kl_per_group(kl_all)
total_kl = torch.sum(kl_coeff_i)
# kl = ( sum * kl / alpha )
kl_coeff_i = kl_coeff_i / alpha_i * total_kl
kl_coeff_i = kl_coeff_i / torch.mean(kl_coeff_i, dim=1, keepdim=True)
kl = torch.sum(kl_all * kl_coeff_i.detach(), dim=1)
# for reporting
kl_coeffs = kl_coeff_i.squeeze(0)
else:
kl_all = torch.stack(kl_all, dim=1)
kl_vals = torch.mean(kl_all, dim=0)
kl = torch.sum(kl_all, dim=1)
kl_coeffs = torch.ones(size=(len(kl_vals),))
return kl_coeff * kl, kl_coeffs, kl_vals
def kl_per_group_vada(all_log_q, all_neg_log_p):
assert(len(all_log_q) == len(all_neg_log_p)
), f'get len={len(all_log_q)} and {len(all_neg_log_p)}'
kl_all_list = []
kl_diag = []
for log_q, neg_log_p in zip(all_log_q, all_neg_log_p):
kl_diag.append(torch.mean(
torch.sum(neg_log_p + log_q, dim=[2, 3]), dim=0))
kl_all_list.append(torch.sum(neg_log_p + log_q,
dim=[1, 2, 3])) # sum over D,H,W
# kl_all = torch.stack(kl_all, dim=1) # batch x num_total_groups
kl_vals = torch.mean(torch.stack(kl_all_list, dim=1),
dim=0) # mean per group
return kl_all_list, kl_vals, kl_diag
def kl_coeff(step, total_step, constant_step, min_kl_coeff, max_kl_coeff):
# return max(min(max_kl_coeff * (step - constant_step) / total_step, max_kl_coeff), min_kl_coeff)
return max(min(min_kl_coeff + (max_kl_coeff - min_kl_coeff) * (step - constant_step) / total_step, max_kl_coeff), min_kl_coeff)
def log_iw(decoder, x, log_q, log_p, crop=False):
recon = reconstruction_loss(decoder, x, crop)
return - recon - log_q + log_p
def reconstruction_loss(decoder, x, crop=False):
recon = decoder.log_p(x)
if crop:
recon = recon[:, :, 2:30, 2:30]
if isinstance(decoder, DiscMixLogistic):
return - torch.sum(recon, dim=[1, 2]) # summation over RGB is done.
else:
return - torch.sum(recon, dim=[1, 2, 3])
def vae_terms(all_log_q, all_eps):
# compute kl
kl_all = []
kl_diag = []
log_p, log_q = 0., 0.
for log_q_conv, eps in zip(all_log_q, all_eps):
log_p_conv = log_p_standard_normal(eps)
kl_per_var = log_q_conv - log_p_conv
kl_diag.append(torch.mean(torch.sum(kl_per_var, dim=[2, 3]), dim=0))
kl_all.append(torch.sum(kl_per_var, dim=[1, 2, 3]))
log_q += torch.sum(log_q_conv, dim=[1, 2, 3])
log_p += torch.sum(log_p_conv, dim=[1, 2, 3])
return log_q, log_p, kl_all, kl_diag
def sum_log_q(all_log_q):
log_q = 0.
for log_q_conv in all_log_q:
log_q += torch.sum(log_q_conv, dim=[1, 2, 3])
return log_q
def cross_entropy_normal(all_eps):
cross_entropy = 0.
neg_log_p_per_group = []
for eps in all_eps:
neg_log_p_conv = - log_p_standard_normal(eps)
neg_log_p = torch.sum(neg_log_p_conv, dim=[1, 2, 3])
cross_entropy += neg_log_p
neg_log_p_per_group.append(neg_log_p_conv)
return cross_entropy, neg_log_p_per_group
def tile_image(batch_image, n, m=None):
if m is None:
m = n
assert n * m == batch_image.size(0)
channels, height, width = batch_image.size(
1), batch_image.size(2), batch_image.size(3)
batch_image = batch_image.view(n, m, channels, height, width)
batch_image = batch_image.permute(2, 0, 3, 1, 4) # n, height, n, width, c
batch_image = batch_image.contiguous().view(channels, n * height, m * width)
return batch_image
def average_gradients_naive(params, is_distributed):
""" Gradient averaging. """
if is_distributed:
size = float(dist.get_world_size())
for param in params:
if param.requires_grad:
param.grad.data /= size
dist.all_reduce(param.grad.data, op=dist.ReduceOp.SUM)
def average_gradients(params, is_distributed):
""" Gradient averaging. """
if is_distributed:
if isinstance(params, types.GeneratorType):
params = [p for p in params]
size = float(dist.get_world_size())
grad_data = []
grad_size = []
grad_shapes = []
# Gather all grad values
for param in params:
if param.requires_grad:
if param.grad is not None:
grad_size.append(param.grad.data.numel())
grad_shapes.append(list(param.grad.data.shape))
grad_data.append(param.grad.data.flatten())
grad_data = torch.cat(grad_data).contiguous()
# All-reduce grad values
grad_data /= size
dist.all_reduce(grad_data, op=dist.ReduceOp.SUM)
# Put back the reduce grad values to parameters
base = 0
i = 0
for param in params:
if param.requires_grad and param.grad is not None:
param.grad.data = grad_data[base:base +
grad_size[i]].view(grad_shapes[i])
base += grad_size[i]
i += 1
def average_params(params, is_distributed):
""" parameter averaging. """
if is_distributed:
size = float(dist.get_world_size())
for param in params:
param.data /= size
dist.all_reduce(param.data, op=dist.ReduceOp.SUM)
def average_tensor(t, is_distributed):
if is_distributed:
size = float(dist.get_world_size())
dist.all_reduce(t.data, op=dist.ReduceOp.SUM)
t.data /= size
def broadcast_params(params, is_distributed):
if is_distributed:
for param in params:
dist.broadcast(param.data, src=0)
def num_output(dataset):
if dataset in {'mnist', 'omniglot'}:
return 28 * 28
elif dataset == 'cifar10':
return 3 * 32 * 32
elif dataset.startswith('celeba') or dataset.startswith('imagenet') or dataset.startswith('lsun'):
size = int(dataset.split('_')[-1])
return 3 * size * size
elif dataset == 'ffhq':
return 3 * 256 * 256
else:
raise NotImplementedError
def get_input_size(dataset):
if dataset in {'mnist', 'omniglot'}:
return 32
elif dataset == 'cifar10':
return 32
elif dataset.startswith('celeba') or dataset.startswith('imagenet') or dataset.startswith('lsun'):
size = int(dataset.split('_')[-1])
return size
elif dataset == 'ffhq':
return 256
elif dataset.startswith('shape'):
return 1 # 2048
else:
raise NotImplementedError
def get_bpd_coeff(dataset):
n = num_output(dataset)
return 1. / np.log(2.) / n
def get_channel_multiplier(dataset, num_scales):
if dataset in {'cifar10', 'omniglot'}:
mult = (1, 1, 1)
elif dataset in {'celeba_256', 'ffhq', 'lsun_church_256'}:
if num_scales == 3:
mult = (1, 1, 1) # used for prior at 16
elif num_scales == 4:
mult = (1, 2, 2, 2) # used for prior at 32
elif num_scales == 5:
mult = (1, 1, 2, 2, 2) # used for prior at 64
elif dataset == 'mnist':
mult = (1, 1)
else:
mult = (1, 1)
# raise NotImplementedError
return mult
def get_attention_scales(dataset):
if dataset in {'cifar10', 'omniglot'}:
attn = (True, False, False)
elif dataset in {'celeba_256', 'ffhq', 'lsun_church_256'}:
# attn = (False, True, False, False) # used for 32
attn = (False, False, True, False, False) # used for 64
elif dataset == 'mnist':
attn = (True, False)
else:
raise NotImplementedError
return attn
def change_bit_length(x, num_bits):
if num_bits != 8:
x = torch.floor(x * 255 / 2 ** (8 - num_bits))
x /= (2 ** num_bits - 1)
return x
def view4D(t, size, inplace=True):
"""
Equal to view(-1, 1, 1, 1).expand(size)
Designed because of this bug:
https://github.com/pytorch/pytorch/pull/48696
"""
if inplace:
return t.unsqueeze_(-1).unsqueeze_(-1).unsqueeze_(-1).expand(size)
else:
return t.unsqueeze(-1).unsqueeze(-1).unsqueeze(-1).expand(size)
def get_arch_cells(arch_type, use_se):
if arch_type == 'res_mbconv':
arch_cells = dict()
arch_cells['normal_enc'] = {'conv_branch': [
'res_bnswish', 'res_bnswish'], 'se': use_se}
arch_cells['down_enc'] = {'conv_branch': [
'res_bnswish', 'res_bnswish'], 'se': use_se}
arch_cells['normal_dec'] = {
'conv_branch': ['mconv_e6k5g0'], 'se': use_se}
arch_cells['up_dec'] = {'conv_branch': ['mconv_e6k5g0'], 'se': use_se}
arch_cells['normal_pre'] = {'conv_branch': [
'res_bnswish', 'res_bnswish'], 'se': use_se}
arch_cells['down_pre'] = {'conv_branch': [
'res_bnswish', 'res_bnswish'], 'se': use_se}
arch_cells['normal_post'] = {
'conv_branch': ['mconv_e3k5g0'], 'se': use_se}
arch_cells['up_post'] = {'conv_branch': ['mconv_e3k5g0'], 'se': use_se}
arch_cells['ar_nn'] = ['']
elif arch_type == 'res_bnswish':
arch_cells = dict()
arch_cells['normal_enc'] = {'conv_branch': [
'res_bnswish', 'res_bnswish'], 'se': use_se}
arch_cells['down_enc'] = {'conv_branch': [
'res_bnswish', 'res_bnswish'], 'se': use_se}
arch_cells['normal_dec'] = {'conv_branch': [
'res_bnswish', 'res_bnswish'], 'se': use_se}
arch_cells['up_dec'] = {'conv_branch': [
'res_bnswish', 'res_bnswish'], 'se': use_se}
arch_cells['normal_pre'] = {'conv_branch': [
'res_bnswish', 'res_bnswish'], 'se': use_se}
arch_cells['down_pre'] = {'conv_branch': [
'res_bnswish', 'res_bnswish'], 'se': use_se}
arch_cells['normal_post'] = {'conv_branch': [
'res_bnswish', 'res_bnswish'], 'se': use_se}
arch_cells['up_post'] = {'conv_branch': [
'res_bnswish', 'res_bnswish'], 'se': use_se}
arch_cells['ar_nn'] = ['']
elif arch_type == 'res_bnswish2':
arch_cells = dict()
arch_cells['normal_enc'] = {
'conv_branch': ['res_bnswish_x2'], 'se': use_se}
arch_cells['down_enc'] = {
'conv_branch': ['res_bnswish_x2'], 'se': use_se}
arch_cells['normal_dec'] = {
'conv_branch': ['res_bnswish_x2'], 'se': use_se}
arch_cells['up_dec'] = {'conv_branch': [
'res_bnswish_x2'], 'se': use_se}
arch_cells['normal_pre'] = {
'conv_branch': ['res_bnswish_x2'], 'se': use_se}
arch_cells['down_pre'] = {
'conv_branch': ['res_bnswish_x2'], 'se': use_se}
arch_cells['normal_post'] = {
'conv_branch': ['res_bnswish_x2'], 'se': use_se}
arch_cells['up_post'] = {'conv_branch': [
'res_bnswish_x2'], 'se': use_se}
arch_cells['ar_nn'] = ['']
elif arch_type == 'res_mbconv_attn':
arch_cells = dict()
arch_cells['normal_enc'] = {'conv_branch': [
'res_bnswish', 'res_bnswish', ], 'se': use_se, 'attn_type': 'attn'}
arch_cells['down_enc'] = {'conv_branch': [
'res_bnswish', 'res_bnswish'], 'se': use_se, 'attn_type': 'attn'}
arch_cells['normal_dec'] = {'conv_branch': [
'mconv_e6k5g0'], 'se': use_se, 'attn_type': 'attn'}
arch_cells['up_dec'] = {'conv_branch': [
'mconv_e6k5g0'], 'se': use_se, 'attn_type': 'attn'}
arch_cells['normal_pre'] = {'conv_branch': [
'res_bnswish', 'res_bnswish'], 'se': use_se}
arch_cells['down_pre'] = {'conv_branch': [
'res_bnswish', 'res_bnswish'], 'se': use_se}
arch_cells['normal_post'] = {
'conv_branch': ['mconv_e3k5g0'], 'se': use_se}
arch_cells['up_post'] = {'conv_branch': ['mconv_e3k5g0'], 'se': use_se}
arch_cells['ar_nn'] = ['']
elif arch_type == 'res_mbconv_attn_half':
arch_cells = dict()
arch_cells['normal_enc'] = {'conv_branch': [
'res_bnswish', 'res_bnswish'], 'se': use_se}
arch_cells['down_enc'] = {'conv_branch': [
'res_bnswish', 'res_bnswish'], 'se': use_se}
arch_cells['normal_dec'] = {'conv_branch': [
'mconv_e6k5g0'], 'se': use_se, 'attn_type': 'attn'}
arch_cells['up_dec'] = {'conv_branch': [
'mconv_e6k5g0'], 'se': use_se, 'attn_type': 'attn'}
arch_cells['normal_pre'] = {'conv_branch': [
'res_bnswish', 'res_bnswish'], 'se': use_se}
arch_cells['down_pre'] = {'conv_branch': [
'res_bnswish', 'res_bnswish'], 'se': use_se}
arch_cells['normal_post'] = {
'conv_branch': ['mconv_e3k5g0'], 'se': use_se}
arch_cells['up_post'] = {'conv_branch': ['mconv_e3k5g0'], 'se': use_se}
arch_cells['ar_nn'] = ['']
else:
raise NotImplementedError
return arch_cells
def get_arch_cells_denoising(arch_type, use_se, apply_sqrt2):
if arch_type == 'res_mbconv':
arch_cells = dict()
arch_cells['normal_enc_diff'] = {
'conv_branch': ['mconv_e6k5g0_gn'], 'se': use_se}
arch_cells['down_enc_diff'] = {
'conv_branch': ['mconv_e6k5g0_gn'], 'se': use_se}
arch_cells['normal_dec_diff'] = {
'conv_branch': ['mconv_e6k5g0_gn'], 'se': use_se}
arch_cells['up_dec_diff'] = {
'conv_branch': ['mconv_e6k5g0_gn'], 'se': use_se}
elif arch_type == 'res_ho':
arch_cells = dict()
arch_cells['normal_enc_diff'] = {
'conv_branch': ['res_gnswish_x2'], 'se': use_se}
arch_cells['down_enc_diff'] = {
'conv_branch': ['res_gnswish_x2'], 'se': use_se}
arch_cells['normal_dec_diff'] = {
'conv_branch': ['res_gnswish_x2'], 'se': use_se}
arch_cells['up_dec_diff'] = {
'conv_branch': ['res_gnswish_x2'], 'se': use_se}
elif arch_type == 'res_ho_p1':
arch_cells = dict()
arch_cells['normal_enc_diff'] = {
'conv_branch': ['res_gnswish_x2_p1'], 'se': use_se}
arch_cells['down_enc_diff'] = {
'conv_branch': ['res_gnswish_x2_p1'], 'se': use_se}
arch_cells['normal_dec_diff'] = {
'conv_branch': ['res_gnswish_x2_p1'], 'se': use_se}
arch_cells['up_dec_diff'] = {
'conv_branch': ['res_gnswish_x2_p1'], 'se': use_se}
elif arch_type == 'res_ho_attn':
arch_cells = dict()
arch_cells['normal_enc_diff'] = {
'conv_branch': ['res_gnswish_x2'], 'se': use_se}
arch_cells['down_enc_diff'] = {
'conv_branch': ['res_gnswish_x2'], 'se': use_se}
arch_cells['normal_dec_diff'] = {
'conv_branch': ['res_gnswish_x2'], 'se': use_se}
arch_cells['up_dec_diff'] = {
'conv_branch': ['res_gnswish_x2'], 'se': use_se}
else:
raise NotImplementedError
for k in arch_cells:
arch_cells[k]['apply_sqrt2'] = apply_sqrt2
return arch_cells
def groups_per_scale(num_scales, num_groups_per_scale):
g = []
n = num_groups_per_scale
for s in range(num_scales):
assert n >= 1
g.append(n)
return g
#class PositionalEmbedding(nn.Module):
# def __init__(self, embedding_dim, scale):
# super(PositionalEmbedding, self).__init__()
# self.embedding_dim = embedding_dim
# self.scale = scale
#
# def forward(self, timesteps):
# assert len(timesteps.shape) == 1
# timesteps = timesteps * self.scale
# half_dim = self.embedding_dim // 2
# emb = math.log(10000) / (half_dim - 1)
# emb = torch.exp(torch.arange(half_dim) * -emb)
# emb = emb.to(device=timesteps.device)
# emb = timesteps[:, None] * emb[None, :]
# emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1)
# return emb
#
#
#class RandomFourierEmbedding(nn.Module):
# def __init__(self, embedding_dim, scale):
# super(RandomFourierEmbedding, self).__init__()
# self.w = nn.Parameter(torch.randn(
# size=(1, embedding_dim // 2)) * scale, requires_grad=False)
#
# def forward(self, timesteps):
# emb = torch.mm(timesteps[:, None], self.w * 2 * 3.14159265359)
# return torch.cat([torch.sin(emb), torch.cos(emb)], dim=1)
#
#
#def init_temb_fun(embedding_type, embedding_scale, embedding_dim):
# if embedding_type == 'positional':
# temb_fun = PositionalEmbedding(embedding_dim, embedding_scale)
# elif embedding_type == 'fourier':
# temb_fun = RandomFourierEmbedding(embedding_dim, embedding_scale)
# else:
# raise NotImplementedError
#
# return temb_fun
def symmetrize_image_data(images):
return 2.0 * images - 1.0
def unsymmetrize_image_data(images):
return (images + 1.) / 2.
def normalize_symmetric(images):
"""
Normalize images by dividing the largest intensity. Used for visualizing the intermediate steps.
"""
b = images.shape[0]
m, _ = torch.max(torch.abs(images).view(b, -1), dim=1)
images /= (m.view(b, 1, 1, 1) + 1e-3)
return images
@torch.jit.script
def soft_clamp5(x: torch.Tensor):
# 5. * torch.tanh(x / 5.) <--> soft differentiable clamp between [-5, 5]
return x.div(5.).tanh_().mul(5.)
@torch.jit.script
def soft_clamp(x: torch.Tensor, a: torch.Tensor):
return x.div(a).tanh_().mul(a)
class SoftClamp5(nn.Module):
def __init__(self):
super(SoftClamp5, self).__init__()
def forward(self, x):
return soft_clamp5(x)
def override_architecture_fields(args, stored_args, logging):
# list of architecture parameters used in NVAE:
architecture_fields = ['arch_instance', 'num_nf', 'num_latent_scales', 'num_groups_per_scale',
'num_latent_per_group', 'num_channels_enc', 'num_preprocess_blocks',
'num_preprocess_cells', 'num_cell_per_cond_enc', 'num_channels_dec',
'num_postprocess_blocks', 'num_postprocess_cells', 'num_cell_per_cond_dec',
'decoder_dist', 'num_x_bits', 'log_sig_q_scale', 'latent_grad_cutoff',
'progressive_output_vae', 'progressive_input_vae', 'channel_mult']
# backward compatibility
""" We have broken backward compatibility. No need to se these manually
if not hasattr(stored_args, 'log_sig_q_scale'):
logging.info('*** Setting %s manually ****', 'log_sig_q_scale')
setattr(stored_args, 'log_sig_q_scale', 5.)
if not hasattr(stored_args, 'latent_grad_cutoff'):
logging.info('*** Setting %s manually ****', 'latent_grad_cutoff')
setattr(stored_args, 'latent_grad_cutoff', 0.)
if not hasattr(stored_args, 'progressive_input_vae'):
logging.info('*** Setting %s manually ****', 'progressive_input_vae')
setattr(stored_args, 'progressive_input_vae', 'none')
if not hasattr(stored_args, 'progressive_output_vae'):
logging.info('*** Setting %s manually ****', 'progressive_output_vae')
setattr(stored_args, 'progressive_output_vae', 'none')
"""
for f in architecture_fields:
if not hasattr(args, f) or getattr(args, f) != getattr(stored_args, f):
logging.info('Setting %s from loaded checkpoint', f)
setattr(args, f, getattr(stored_args, f))
def init_processes(rank, size, fn, args, config):
""" Initialize the distributed environment. """
os.environ['MASTER_ADDR'] = args.master_address
os.environ['MASTER_PORT'] = '6020'
if args.num_proc_node == 1:
import socket
import errno
a_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
for p in range(6010, 6030):
location = (args.master_address, p) # "127.0.0.1", p)
try:
a_socket.bind((args.master_address, p))
logger.debug('set port as {}', p)
os.environ['MASTER_PORT'] = '%d' % p
a_socket.close()
break
except socket.error as e:
a = 0
# if e.errno == errno.EADDRINUSE:
# # logger.debug("Port {} is already in use", p)
# else:
# logger.debug(e)
logger.info('init_process: rank={}, world_size={}', rank, size)
torch.cuda.set_device(args.local_rank)
dist.init_process_group(
backend='nccl', init_method='env://', rank=rank, world_size=size)
fn(args, config)
logger.info('barrier: rank={}, world_size={}', rank, size)
dist.barrier()
logger.info('skip destroy_process_group: rank={}, world_size={}', rank, size)
# dist.destroy_process_group()
logger.info('skip destroy fini')
def sample_rademacher_like(y):
return torch.randint(low=0, high=2, size=y.shape, device='cuda') * 2 - 1
def sample_gaussian_like(y):
return torch.randn_like(y, device='cuda')
def trace_df_dx_hutchinson(f, x, noise, no_autograd):
"""
Hutchinson's trace estimator for Jacobian df/dx, O(1) call to autograd
"""
if no_autograd:
# the following is compatible with checkpointing
torch.sum(f * noise).backward()
# torch.autograd.backward(tensors=[f], grad_tensors=[noise])
jvp = x.grad
trJ = torch.sum(jvp * noise, dim=[1, 2, 3])
x.grad = None
else:
jvp = torch.autograd.grad(f, x, noise, create_graph=False)[0]
trJ = torch.sum(jvp * noise, dim=[1, 2, 3])
# trJ = torch.einsum('bijk,bijk->b', jvp, noise) # we could test if there's a speed difference in einsum vs sum
return trJ
def calc_jacobian_regularization(pred_params, eps_t, dae, var_t, m_t, f_t, g2_t, var_N_t, args):
"""
Calculates Jabobian regularization loss. For reference implementations, see
https://github.com/facebookresearch/jacobian_regularizer/blob/master/jacobian/jacobian.py or
https://github.com/cfinlay/ffjord-rnode/blob/master/lib/layers/odefunc.py.
"""
# eps_t_jvp = eps_t.detach()
# eps_t_jvp = eps_t.detach().requires_grad_()
if args.no_autograd_jvp:
raise NotImplementedError(
"We have not implemented no_autograd_jvp for jacobian reg.")
jvp_ode_func_norms = []
alpha = torch.sigmoid(dae.mixing_logit.detach())
for _ in range(args.jac_reg_samples):
noise = sample_gaussian_like(eps_t)
jvp = torch.autograd.grad(
pred_params, eps_t, noise, create_graph=True)[0]
if args.sde_type in ['geometric_sde', 'vpsde', 'power_vpsde']:
jvp_ode_func = alpha * (noise * torch.sqrt(var_t) - jvp)
if not args.jac_kin_reg_drop_weights:
jvp_ode_func = f_t / torch.sqrt(var_t) * jvp_ode_func
elif args.sde_type in ['sub_vpsde', 'sub_power_vpsde']:
sigma2_N_t = (1.0 - m_t ** 2) ** 2 + m_t ** 2
jvp_ode_func = noise * torch.sqrt(var_t) / (1.0 - m_t ** 4) - (
(1.0 - alpha) * noise * torch.sqrt(var_t) / sigma2_N_t + alpha * jvp)
if not args.jac_kin_reg_drop_weights:
jvp_ode_func = f_t * (1.0 - m_t ** 4) / \
torch.sqrt(var_t) * jvp_ode_func
elif args.sde_type in ['vesde']:
jvp_ode_func = (1.0 - alpha) * noise * \
torch.sqrt(var_t) / var_N_t + alpha * jvp
if not args.jac_kin_reg_drop_weights:
jvp_ode_func = 0.5 * g2_t / torch.sqrt(var_t) * jvp_ode_func
else:
raise ValueError("Unrecognized SDE type: {}".format(args.sde_type))
jvp_ode_func_norms.append(jvp_ode_func.view(
eps_t.size(0), -1).pow(2).sum(dim=1, keepdim=True))
jac_reg_loss = torch.cat(jvp_ode_func_norms, dim=1).mean()
# jac_reg_loss = torch.mean(jvp_ode_func.view(eps_t.size(0), -1).pow(2).sum(dim=1))
return jac_reg_loss
def calc_kinetic_regularization(pred_params, eps_t, dae, var_t, m_t, f_t, g2_t, var_N_t, args):
"""
Calculates kinetic regularization loss. For a reference implementation, see
https://github.com/cfinlay/ffjord-rnode/blob/master/lib/layers/wrappers/cnf_regularization.py
"""
# eps_t_kin = eps_t.detach()
alpha = torch.sigmoid(dae.mixing_logit.detach())
if args.sde_type in ['geometric_sde', 'vpsde', 'power_vpsde']:
ode_func = alpha * (eps_t * torch.sqrt(var_t) - pred_params)
if not args.jac_kin_reg_drop_weights:
ode_func = f_t / torch.sqrt(var_t) * ode_func
elif args.sde_type in ['sub_vpsde', 'sub_power_vpsde']:
sigma2_N_t = (1.0 - m_t ** 2) ** 2 + m_t ** 2
ode_func = eps_t * torch.sqrt(var_t) / (1.0 - m_t ** 4) - (
(1.0 - alpha) * eps_t * torch.sqrt(var_t) / sigma2_N_t + alpha * pred_params)
if not args.jac_kin_reg_drop_weights:
ode_func = f_t * (1.0 - m_t ** 4) / torch.sqrt(var_t) * ode_func
elif args.sde_type in ['vesde']:
ode_func = (1.0 - alpha) * eps_t * torch.sqrt(var_t) / \
var_N_t + alpha * pred_params
if not args.jac_kin_reg_drop_weights:
ode_func = 0.5 * g2_t / torch.sqrt(var_t) * ode_func
else:
raise ValueError("Unrecognized SDE type: {}".format(args.sde_type))
kin_reg_loss = torch.mean(ode_func.view(
eps_t.size(0), -1).pow(2).sum(dim=1))
return kin_reg_loss
def different_p_q_objectives(iw_sample_p, iw_sample_q):
assert iw_sample_p in ['ll_uniform', 'drop_all_uniform', 'll_iw', 'drop_all_iw', 'drop_sigma2t_iw', 'rescale_iw',
'drop_sigma2t_uniform']
assert iw_sample_q in ['reweight_p_samples', 'll_uniform', 'll_iw']
# Removed assert below. It may be stupid, but user can still do it. It may make sense for debugging purposes.
# assert iw_sample_p != iw_sample_q, 'It does not make sense to use the same objectives for p and q, but train ' \
# 'with separated q and p updates. To reuse the p objective for q, specify ' \
# '"reweight_p_samples" instead (for the ll-based objectives, the ' \
# 'reweighting factor will simply be 1.0 then)!'
# In these cases, we reuse the likelihood-based p-objective (either the uniform sampling version or the importance
# sampling version) also for q.
if iw_sample_p in ['ll_uniform', 'll_iw'] and iw_sample_q == 'reweight_p_samples':
return False
# In these cases, we are using a non-likelihood-based objective for p, and hence definitly need to use another q
# objective.
else:
return True
def decoder_output(dataset, logits, fixed_log_scales=None):
if dataset in {'cifar10', 'celeba_64', 'celeba_256', 'imagenet_32', 'imagenet_64', 'ffhq',
'lsun_bedroom_128', 'lsun_bedroom_256', 'mnist', 'omniglot',
'lsun_church_256'}:
return PixelNormal(logits, fixed_log_scales)
else:
return PixelNormal(logits, fixed_log_scales)
# raise NotImplementedError
def get_mixed_prediction(mixed_prediction, param, mixing_logit, mixing_component=None):
if mixed_prediction:
assert mixing_component is not None, 'Provide mixing component when mixed_prediction is enabled.'
coeff = torch.sigmoid(mixing_logit)
param = (1 - coeff) * mixing_component + coeff * param
return param
def set_vesde_sigma_max(args, vae, train_queue, logging, is_distributed):
logging.info('')
logging.info(
'Calculating max. pairwise distance in latent space to set sigma2_max for VESDE...')
eps_list = []
vae.eval()
for step, x in enumerate(train_queue):
x = x[0] if len(x) > 1 else x
x = x.cuda()
x = symmetrize_image_data(x)
# run vae
with autocast(enabled=args.autocast_train):
with torch.set_grad_enabled(False):
logits, all_log_q, all_eps = vae(x)
eps = torch.cat(all_eps, dim=1)
eps_list.append(eps.detach())
# if step > 5: ### DEBUG
# break ### DEBUG
# concat eps tensor on each GPU and then gather all on all GPUs
eps_this_rank = torch.cat(eps_list, dim=0)
if is_distributed:
eps_all_gathered = [torch.zeros_like(
eps_this_rank)] * dist.get_world_size()
dist.all_gather(eps_all_gathered, eps_this_rank)
eps_full = torch.cat(eps_all_gathered, dim=0)
else:
eps_full = eps_this_rank
# max pairwise distance squared between all latent encodings, is computed on CPU
eps_full = eps_full.cpu().float()
eps_full = eps_full.flatten(start_dim=1).unsqueeze(0)
max_pairwise_dist_sqr = torch.cdist(eps_full, eps_full).square().max()
max_pairwise_dist_sqr = max_pairwise_dist_sqr.cuda()
# to be safe, we broadcast to all GPUs if we are in distributed environment. Shouldn't be necessary in principle.
if is_distributed:
dist.broadcast(max_pairwise_dist_sqr, src=0)
args.sigma2_max = max_pairwise_dist_sqr.item()
logging.info('Done! Set args.sigma2_max set to {}'.format(args.sigma2_max))
logging.info('')
return args
def mask_inactive_variables(x, is_active):
x = x * is_active
return x
def common_x_operations(x, num_x_bits):
x = x[0] if len(x) > 1 else x
x = x.cuda()
# change bit length
x = change_bit_length(x, num_x_bits)
x = symmetrize_image_data(x)
return x
def vae_regularization(args, vae_sn_calculator, loss_weight=None):
"""
when using hvae_trainer, we pass args=None, and loss_weight value
"""
regularization_q, vae_norm_loss, vae_bn_loss, vae_wdn_coeff = 0., 0., 0., args.weight_decay_norm_vae if loss_weight is None else loss_weight
if loss_weight is not None or args.train_vae:
vae_norm_loss = vae_sn_calculator.spectral_norm_parallel()
vae_bn_loss = vae_sn_calculator.batchnorm_loss()
regularization_q = (vae_norm_loss + vae_bn_loss) * vae_wdn_coeff
return regularization_q, vae_norm_loss, vae_bn_loss, vae_wdn_coeff
def dae_regularization(args, dae_sn_calculator, diffusion, dae, step, t, pred_params_p, eps_t_p, var_t_p, m_t_p, g2_t_p):
dae_wdn_coeff = args.weight_decay_norm_dae
dae_norm_loss = dae_sn_calculator.spectral_norm_parallel()
dae_bn_loss = dae_sn_calculator.batchnorm_loss()
regularization_p = (dae_norm_loss + dae_bn_loss) * dae_wdn_coeff
# Jacobian regularization
jac_reg_loss = 0.
if args.jac_reg_coeff > 0.0 and step % args.jac_reg_freq == 0:
f_t = diffusion.f(t).view(-1, 1, 1, 1)
var_N_t = diffusion.var_N(
t).view(-1, 1, 1, 1) if args.sde_type == 'vesde' else None
"""
# Arash: Please remove the following if it looks correct to you, Karsten.
# jac_reg_loss = utils.calc_jacobian_regularization(pred_params, eps_t, dae, var_t, m_t, f_t, args)
if args.iw_sample_q in ['ll_uniform', 'll_iw']:
pred_params_jac_reg = torch.chunk(pred_params, chunks=2, dim=0)[0]
var_t_jac_reg, m_t_jac_reg, f_t_jac_reg = torch.chunk(var_t, chunks=2, dim=0)[0], \
torch.chunk(m_t, chunks=2, dim=0)[0], \
torch.chunk(f_t, chunks=2, dim=0)[0]
g2_t_jac_reg = torch.chunk(g2_t, chunks=2, dim=0)[0]
var_N_t_jac_reg = torch.chunk(var_N_t, chunks=2, dim=0)[0] if args.sde_type == 'vesde' else None
else:
pred_params_jac_reg = pred_params
var_t_jac_reg, m_t_jac_reg, f_t_jac_reg, g2_t_jac_reg, var_N_t_jac_reg = var_t, m_t, f_t, g2_t, var_N_t
jac_reg_loss = utils.calc_jacobian_regularization(pred_params_jac_reg, eps_t_p, dae, var_t_jac_reg, m_t_jac_reg,
f_t_jac_reg, g2_t_jac_reg, var_N_t_jac_reg, args)
"""
jac_reg_loss = calc_jacobian_regularization(pred_params_p, eps_t_p, dae, var_t_p, m_t_p,
f_t, g2_t_p, var_N_t, args)
regularization_p += args.jac_reg_coeff * jac_reg_loss
# Kinetic regularization
kin_reg_loss = 0.
if args.kin_reg_coeff > 0.0:
f_t = diffusion.f(t).view(-1, 1, 1, 1)
var_N_t = diffusion.var_N(
t).view(-1, 1, 1, 1) if args.sde_type == 'vesde' else None
"""
# Arash: Please remove the following if it looks correct to you, Karsten.
# kin_reg_loss = utils.calc_kinetic_regularization(pred_params, eps_t, dae, var_t, m_t, f_t, args)
if args.iw_sample_q in ['ll_uniform', 'll_iw']:
pred_params_kin_reg = torch.chunk(pred_params, chunks=2, dim=0)[0]
var_t_kin_reg, m_t_kin_reg, f_t_kin_reg = torch.chunk(var_t, chunks=2, dim=0)[0], \
torch.chunk(m_t, chunks=2, dim=0)[0], \
torch.chunk(f_t, chunks=2, dim=0)[0]
g2_t_kin_reg = torch.chunk(g2_t, chunks=2, dim=0)[0]
var_N_t_kin_reg = torch.chunk(var_N_t, chunks=2, dim=0)[0] if args.sde_type == 'vesde' else None
else:
pred_params_kin_reg = pred_params
var_t_kin_reg, m_t_kin_reg, f_t_kin_reg, g2_t_kin_reg, var_N_t_kin_reg = var_t, m_t, f_t, g2_t, var_N_t
kin_reg_loss = utils.calc_kinetic_regularization(pred_params_kin_reg, eps_t_p, dae, var_t_kin_reg, m_t_kin_reg,
f_t_kin_reg, g2_t_kin_reg, var_N_t_kin_reg, args)
"""
kin_reg_loss = calc_kinetic_regularization(pred_params_p, eps_t_p, dae, var_t_p, m_t_p,
f_t, g2_t_p, var_N_t, args)
regularization_p += args.kin_reg_coeff * kin_reg_loss
return regularization_p, dae_norm_loss, dae_bn_loss, dae_wdn_coeff, jac_reg_loss, kin_reg_loss
def update_vae_lr(args, global_step, warmup_iters, vae_optimizer):
if global_step < warmup_iters:
lr = args.trainer.opt.lr * float(global_step) / warmup_iters
for param_group in vae_optimizer.param_groups:
param_group['lr'] = lr
# use same lr if lr for local-dae is not specified
def update_lr(args, global_step, warmup_iters, dae_optimizer, vae_optimizer, dae_local_optimizer=None):
if global_step < warmup_iters:
lr = args.learning_rate_dae * float(global_step) / warmup_iters
if args.learning_rate_mlogit > 0 and len(dae_optimizer.param_groups) > 1:
lr_mlogit = args.learning_rate_mlogit * \
float(global_step) / warmup_iters
for i, param_group in enumerate(dae_optimizer.param_groups):
if i == 0:
param_group['lr'] = lr_mlogit
else:
param_group['lr'] = lr
else:
for param_group in dae_optimizer.param_groups:
param_group['lr'] = lr
# use same lr if lr for local-dae is not specified
lr = lr if args.learning_rate_dae_local <= 0 else args.learning_rate_dae_local * \
float(global_step) / warmup_iters
if dae_local_optimizer is not None:
for param_group in dae_local_optimizer.param_groups:
param_group['lr'] = lr
if args.train_vae:
lr = args.learning_rate_vae * float(global_step) / warmup_iters
for param_group in vae_optimizer.param_groups:
param_group['lr'] = lr
def start_meters():
tr_loss_meter = AvgrageMeter()
vae_recon_meter = AvgrageMeter()
vae_kl_meter = AvgrageMeter()
vae_nelbo_meter = AvgrageMeter()
kl_per_group_ema = AvgrageMeter()
return tr_loss_meter, vae_recon_meter, vae_kl_meter, vae_nelbo_meter, kl_per_group_ema
def epoch_logging(args, writer, step, vae_recon_meter, vae_kl_meter, vae_nelbo_meter, tr_loss_meter, kl_per_group_ema):
average_tensor(vae_recon_meter.avg, args.distributed)
average_tensor(vae_kl_meter.avg, args.distributed)
average_tensor(vae_nelbo_meter.avg, args.distributed)
average_tensor(tr_loss_meter.avg, args.distributed)
average_tensor(kl_per_group_ema.avg, args.distributed)
writer.add_scalar('epoch/vae_recon', vae_recon_meter.avg, step)
writer.add_scalar('epoch/vae_kl', vae_kl_meter.avg, step)
writer.add_scalar('epoch/vae_nelbo', vae_nelbo_meter.avg, step)
writer.add_scalar('epoch/total_loss', tr_loss_meter.avg, step)
# add kl value per group to tensorboard
for i in range(len(kl_per_group_ema.avg)):
writer.add_scalar('kl_value/group_%d' %
i, kl_per_group_ema.avg[i], step)
def infer_active_variables(train_queue, vae, args, device, distributed, max_iter=None):
kl_meter = AvgrageMeter()
vae.eval()
for step, x in enumerate(train_queue):
if max_iter is not None and step > max_iter:
break
tr_pts = x['tr_points']
with autocast(enabled=args.autocast_train):
# apply vae:
with torch.set_grad_enabled(False):
# output = model.recont(val_x) ## torch.cat([val_x, tr_x]))
dist = vae.encode(tr_pts.to(device))
eps = dist.sample()[0]
all_log_q = [dist.log_p(eps)]
## _, all_log_q, all_eps = vae(x)
## all_eps = vae.concat_eps_per_scale(all_eps)
## all_log_q = vae.concat_eps_per_scale(all_log_q)
all_eps = [eps]
def make_4d(xlist): return [
x.unsqueeze(-1).unsqueeze(-1) if len(x.shape) == 2 else x.unsqueeze(-1) for x in xlist]
log_q, log_p, kl_all, kl_diag = vae_terms(
make_4d(all_log_q), make_4d(all_eps))
kl_meter.update(kl_diag[0], 1) # only the top scale
average_tensor(kl_meter.avg, distributed)
return kl_meter.avg > 0.1