import math import sys from collections import OrderedDict sys.path.append('..') import lpips import torch import torch.nn.functional as F from torchvision.utils import save_image from models.archs.vqgan_arch import (Decoder, DecoderRes, Discriminator, Encoder, VectorQuantizerSpatialTextureAware, VectorQuantizerTexture) from models.losses.vqgan_loss import (DiffAugment, adopt_weight, calculate_adaptive_weight, hinge_d_loss) class HierarchyVQSpatialTextureAwareModel(): def __init__(self, opt): self.opt = opt self.device = torch.device('cuda') self.top_encoder = Encoder( ch=opt['top_ch'], num_res_blocks=opt['top_num_res_blocks'], attn_resolutions=opt['top_attn_resolutions'], ch_mult=opt['top_ch_mult'], in_channels=opt['top_in_channels'], resolution=opt['top_resolution'], z_channels=opt['top_z_channels'], double_z=opt['top_double_z'], dropout=opt['top_dropout']).to(self.device) self.decoder = Decoder( in_channels=opt['top_in_channels'], resolution=opt['top_resolution'], z_channels=opt['top_z_channels'], ch=opt['top_ch'], out_ch=opt['top_out_ch'], num_res_blocks=opt['top_num_res_blocks'], attn_resolutions=opt['top_attn_resolutions'], ch_mult=opt['top_ch_mult'], dropout=opt['top_dropout'], resamp_with_conv=True, give_pre_end=False).to(self.device) self.top_quantize = VectorQuantizerTexture( 1024, opt['embed_dim'], beta=0.25).to(self.device) self.top_quant_conv = torch.nn.Conv2d(opt["top_z_channels"], opt['embed_dim'], 1).to(self.device) self.top_post_quant_conv = torch.nn.Conv2d(opt['embed_dim'], opt["top_z_channels"], 1).to(self.device) self.load_top_pretrain_models() self.bot_encoder = Encoder( ch=opt['bot_ch'], num_res_blocks=opt['bot_num_res_blocks'], attn_resolutions=opt['bot_attn_resolutions'], ch_mult=opt['bot_ch_mult'], in_channels=opt['bot_in_channels'], resolution=opt['bot_resolution'], z_channels=opt['bot_z_channels'], double_z=opt['bot_double_z'], dropout=opt['bot_dropout']).to(self.device) self.bot_decoder_res = DecoderRes( in_channels=opt['bot_in_channels'], resolution=opt['bot_resolution'], z_channels=opt['bot_z_channels'], ch=opt['bot_ch'], num_res_blocks=opt['bot_num_res_blocks'], ch_mult=opt['bot_ch_mult'], dropout=opt['bot_dropout'], give_pre_end=False).to(self.device) self.bot_quantize = VectorQuantizerSpatialTextureAware( opt['bot_n_embed'], opt['embed_dim'], beta=0.25, spatial_size=opt['codebook_spatial_size']).to(self.device) self.bot_quant_conv = torch.nn.Conv2d(opt["bot_z_channels"], opt['embed_dim'], 1).to(self.device) self.bot_post_quant_conv = torch.nn.Conv2d(opt['embed_dim'], opt["bot_z_channels"], 1).to(self.device) self.disc = Discriminator( opt['n_channels'], opt['ndf'], n_layers=opt['disc_layers']).to(self.device) self.perceptual = lpips.LPIPS(net="vgg").to(self.device) self.perceptual_weight = opt['perceptual_weight'] self.disc_start_step = opt['disc_start_step'] self.disc_weight_max = opt['disc_weight_max'] self.diff_aug = opt['diff_aug'] self.policy = "color,translation" self.load_discriminator_models() self.disc.train() self.fix_decoder = opt['fix_decoder'] self.init_training_settings() def load_top_pretrain_models(self): # load pretrained vqgan for segmentation mask top_vae_checkpoint = torch.load(self.opt['top_vae_path']) self.top_encoder.load_state_dict( top_vae_checkpoint['encoder'], strict=True) self.decoder.load_state_dict( top_vae_checkpoint['decoder'], strict=True) self.top_quantize.load_state_dict( top_vae_checkpoint['quantize'], strict=True) self.top_quant_conv.load_state_dict( top_vae_checkpoint['quant_conv'], strict=True) self.top_post_quant_conv.load_state_dict( top_vae_checkpoint['post_quant_conv'], strict=True) self.top_encoder.eval() self.top_quantize.eval() self.top_quant_conv.eval() self.top_post_quant_conv.eval() def init_training_settings(self): self.log_dict = OrderedDict() self.configure_optimizers() def configure_optimizers(self): optim_params = [] for v in self.bot_encoder.parameters(): if v.requires_grad: optim_params.append(v) for v in self.bot_decoder_res.parameters(): if v.requires_grad: optim_params.append(v) for v in self.bot_quantize.parameters(): if v.requires_grad: optim_params.append(v) for v in self.bot_quant_conv.parameters(): if v.requires_grad: optim_params.append(v) for v in self.bot_post_quant_conv.parameters(): if v.requires_grad: optim_params.append(v) if not self.fix_decoder: for name, v in self.decoder.named_parameters(): if v.requires_grad: if 'up.0' in name: optim_params.append(v) if 'up.1' in name: optim_params.append(v) if 'up.2' in name: optim_params.append(v) if 'up.3' in name: optim_params.append(v) self.optimizer = torch.optim.Adam(optim_params, lr=self.opt['lr']) self.disc_optimizer = torch.optim.Adam( self.disc.parameters(), lr=self.opt['lr']) def load_discriminator_models(self): # load pretrained vqgan for segmentation mask top_vae_checkpoint = torch.load(self.opt['top_vae_path']) self.disc.load_state_dict( top_vae_checkpoint['discriminator'], strict=True) def save_network(self, save_path): """Save networks. """ save_dict = {} save_dict['bot_encoder'] = self.bot_encoder.state_dict() save_dict['bot_decoder_res'] = self.bot_decoder_res.state_dict() save_dict['decoder'] = self.decoder.state_dict() save_dict['bot_quantize'] = self.bot_quantize.state_dict() save_dict['bot_quant_conv'] = self.bot_quant_conv.state_dict() save_dict['bot_post_quant_conv'] = self.bot_post_quant_conv.state_dict( ) save_dict['discriminator'] = self.disc.state_dict() torch.save(save_dict, save_path) def load_network(self): checkpoint = torch.load(self.opt['pretrained_models']) self.bot_encoder.load_state_dict( checkpoint['bot_encoder'], strict=True) self.bot_decoder_res.load_state_dict( checkpoint['bot_decoder_res'], strict=True) self.decoder.load_state_dict(checkpoint['decoder'], strict=True) self.bot_quantize.load_state_dict( checkpoint['bot_quantize'], strict=True) self.bot_quant_conv.load_state_dict( checkpoint['bot_quant_conv'], strict=True) self.bot_post_quant_conv.load_state_dict( checkpoint['bot_post_quant_conv'], strict=True) def optimize_parameters(self, data, step): self.bot_encoder.train() self.bot_decoder_res.train() if not self.fix_decoder: self.decoder.train() self.bot_quantize.train() self.bot_quant_conv.train() self.bot_post_quant_conv.train() loss, d_loss = self.training_step(data, step) self.optimizer.zero_grad() loss.backward() self.optimizer.step() if step > self.disc_start_step: self.disc_optimizer.zero_grad() d_loss.backward() self.disc_optimizer.step() def top_encode(self, x, mask): h = self.top_encoder(x) h = self.top_quant_conv(h) quant, _, _ = self.top_quantize(h, mask) quant = self.top_post_quant_conv(quant) return quant def bot_encode(self, x, mask): h = self.bot_encoder(x) h = self.bot_quant_conv(h) quant, emb_loss, info = self.bot_quantize(h, mask) quant = self.bot_post_quant_conv(quant) bot_dec_res = self.bot_decoder_res(quant) return bot_dec_res, emb_loss, info def decode(self, quant_top, bot_dec_res): dec = self.decoder(quant_top, bot_h=bot_dec_res) return dec def forward_step(self, input, mask): with torch.no_grad(): quant_top = self.top_encode(input, mask) bot_dec_res, diff, _ = self.bot_encode(input, mask) dec = self.decode(quant_top, bot_dec_res) return dec, diff def feed_data(self, data): x = data['image'].float().to(self.device) mask = data['texture_mask'].float().to(self.device) return x, mask def training_step(self, data, step): x, mask = self.feed_data(data) xrec, codebook_loss = self.forward_step(x, mask) # get recon/perceptual loss recon_loss = torch.abs(x.contiguous() - xrec.contiguous()) p_loss = self.perceptual(x.contiguous(), xrec.contiguous()) nll_loss = recon_loss + self.perceptual_weight * p_loss nll_loss = torch.mean(nll_loss) # augment for input to discriminator if self.diff_aug: xrec = DiffAugment(xrec, policy=self.policy) # update generator logits_fake = self.disc(xrec) g_loss = -torch.mean(logits_fake) last_layer = self.decoder.conv_out.weight d_weight = calculate_adaptive_weight(nll_loss, g_loss, last_layer, self.disc_weight_max) d_weight *= adopt_weight(1, step, self.disc_start_step) loss = nll_loss + d_weight * g_loss + codebook_loss self.log_dict["loss"] = loss self.log_dict["l1"] = recon_loss.mean().item() self.log_dict["perceptual"] = p_loss.mean().item() self.log_dict["nll_loss"] = nll_loss.item() self.log_dict["g_loss"] = g_loss.item() self.log_dict["d_weight"] = d_weight self.log_dict["codebook_loss"] = codebook_loss.item() if step > self.disc_start_step: if self.diff_aug: logits_real = self.disc( DiffAugment(x.contiguous().detach(), policy=self.policy)) else: logits_real = self.disc(x.contiguous().detach()) logits_fake = self.disc(xrec.contiguous().detach( )) # detach so that generator isn"t also updated d_loss = hinge_d_loss(logits_real, logits_fake) self.log_dict["d_loss"] = d_loss else: d_loss = None return loss, d_loss @torch.no_grad() def inference(self, data_loader, save_dir): self.bot_encoder.eval() self.bot_decoder_res.eval() self.decoder.eval() self.bot_quantize.eval() self.bot_quant_conv.eval() self.bot_post_quant_conv.eval() loss_total = 0 num = 0 for _, data in enumerate(data_loader): img_name = data['img_name'][0] x, mask = self.feed_data(data) xrec, _ = self.forward_step(x, mask) recon_loss = torch.abs(x.contiguous() - xrec.contiguous()) p_loss = self.perceptual(x.contiguous(), xrec.contiguous()) nll_loss = recon_loss + self.perceptual_weight * p_loss nll_loss = torch.mean(nll_loss) loss_total += nll_loss num += x.size(0) if x.shape[1] > 3: # colorize with random projection assert xrec.shape[1] > 3 # convert logits to indices xrec = torch.argmax(xrec, dim=1, keepdim=True) xrec = F.one_hot(xrec, num_classes=x.shape[1]) xrec = xrec.squeeze(1).permute(0, 3, 1, 2).float() x = self.to_rgb(x) xrec = self.to_rgb(xrec) img_cat = torch.cat([x, xrec], dim=3).detach() img_cat = ((img_cat + 1) / 2) img_cat = img_cat.clamp_(0, 1) save_image( img_cat, f'{save_dir}/{img_name}.png', nrow=1, padding=4) return (loss_total / num).item() def get_current_log(self): return self.log_dict def update_learning_rate(self, epoch): """Update learning rate. Args: current_iter (int): Current iteration. warmup_iter (int): Warmup iter numbers. -1 for no warmup. Default: -1. """ lr = self.optimizer.param_groups[0]['lr'] if self.opt['lr_decay'] == 'step': lr = self.opt['lr'] * ( self.opt['gamma']**(epoch // self.opt['step'])) elif self.opt['lr_decay'] == 'cos': lr = self.opt['lr'] * ( 1 + math.cos(math.pi * epoch / self.opt['num_epochs'])) / 2 elif self.opt['lr_decay'] == 'linear': lr = self.opt['lr'] * (1 - epoch / self.opt['num_epochs']) elif self.opt['lr_decay'] == 'linear2exp': if epoch < self.opt['turning_point'] + 1: # learning rate decay as 95% # at the turning point (1 / 95% = 1.0526) lr = self.opt['lr'] * ( 1 - epoch / int(self.opt['turning_point'] * 1.0526)) else: lr *= self.opt['gamma'] elif self.opt['lr_decay'] == 'schedule': if epoch in self.opt['schedule']: lr *= self.opt['gamma'] else: raise ValueError('Unknown lr mode {}'.format(self.opt['lr_decay'])) # set learning rate for param_group in self.optimizer.param_groups: param_group['lr'] = lr return lr