""" Modification of Open-MAGVIT2 code, including adding gradient accumulation during training, using VQConfig, removing hardcoded arguments and removing unnecessary code. """ import torch import torch.nn.functional as F import lightning as L from collections import OrderedDict from contextlib import contextmanager from magvit2.config import VQConfig from magvit2.modules.diffusionmodules.improved_model import Encoder, Decoder from magvit2.modules.losses.vqperceptual import VQLPIPSWithDiscriminator from magvit2.modules.vqvae.lookup_free_quantize import LFQ from magvit2.modules.scheduler.lr_scheduler import Scheduler_LinearWarmup, Scheduler_LinearWarmup_CosineDecay from magvit2.modules.ema import LitEma class VQModel(L.LightningModule): def __init__( self, config: VQConfig, training_args=None, ckpt_path=None, ignore_keys=[], image_key="image", colorize_nlabels=None, monitor=None, use_ema=True, stage=None, ): super().__init__() self.training_args = training_args self.image_key = image_key self.encoder = Encoder(config) self.decoder = Decoder(config) self.loss = VQLPIPSWithDiscriminator(config) self.quantize = LFQ(config) self.use_ema = use_ema self.stage = stage if ckpt_path is not None: self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys, stage=stage) self.image_key = image_key if colorize_nlabels is not None: assert isinstance(colorize_nlabels, int) self.register_buffer("colorize", torch.randn(3, colorize_nlabels, 1, 1)) if monitor is not None: self.monitor = monitor self.generator_params = list(self.encoder.parameters()) + \ list(self.decoder.parameters()) + \ list(self.quantize.parameters()) if self.use_ema and stage is None: #no need to construct ema when training transformer self.model_ema = LitEma(self) # Note: this means EMA weights are overriden after `init_from_ckpt`. self.automatic_optimization = False self.strict_loading = False @contextmanager def ema_scope(self, context=None): if self.use_ema: self.model_ema.store(self.parameters()) self.model_ema.copy_to(self) if context is not None: print(f"{context}: Switched to EMA weights") try: yield None finally: if self.use_ema: self.model_ema.restore(self.parameters()) if context is not None: print(f"{context}: Restored training weights") def state_dict(self, *args, destination=None, prefix='', keep_vars=False): """ save the state_dict and filter out the """ return {k: v for k, v in super().state_dict(*args, destination, prefix, keep_vars).items() if ("inception_model" not in k and "lpips_vgg" not in k and "lpips_alex" not in k)} def init_from_ckpt(self, path, ignore_keys=list(), stage=None): sd = torch.load(path, map_location="cpu")["state_dict"] ema_mapping = {} new_params = OrderedDict() if stage == "transformer": ### directly use ema encoder and decoder parameter if self.use_ema: for k, v in sd.items(): if "encoder" in k: if "model_ema" in k: k = k.replace("model_ema.", "") #load EMA Encoder or Decoder new_k = ema_mapping[k] new_params[new_k] = v s_name = k.replace('.', '') ema_mapping.update({s_name: k}) continue if "decoder" in k: if "model_ema" in k: k = k.replace("model_ema.", "") # load EMA Encoder or Decoder new_k = ema_mapping[k] new_params[new_k] = v s_name = k.replace(".", "") ema_mapping.update({s_name: k}) continue else: # also only load the Generator for k, v in sd.items(): if "encoder" in k: new_params[k] = v elif "decoder" in k: new_params[k] = v missing_keys, unexpected_keys = self.load_state_dict(new_params, strict=False) else: ## simple resume missing_keys, unexpected_keys = self.load_state_dict(sd, strict=False) # print(f"{missing_keys=} {unexpected_keys=}") print(f"Restored from {path}") def encode_without_quantize(self, x): h = self.encoder(x) return h def encode(self, x, **kwargs): h = self.encoder(x) (quant, emb_loss, info), loss_breakdown = self.quantize(h, return_loss_breakdown=True, **kwargs) ### using token factorization the info is a tuple (each for embedding) return quant, emb_loss, info, loss_breakdown def decode(self, quant): dec = self.decoder(quant) return dec def forward(self, input): quant, codebook_loss, _, loss_break = self.encode(input) dec = self.decode(quant) return dec, codebook_loss, loss_break def get_input(self, batch, image_key): x = batch[image_key] if len(x.shape) == 3: # grayscale case I think? - Kevin x = x[..., None] x = x.permute(0, 3, 1, 2).to(memory_format=torch.contiguous_format) return x.float() # fix mulitple optimizer bug # refer to https://lightning.ai/docs/pytorch/stable/model/manual_optimization.html def training_step(self, batch, batch_idx): x = self.get_input(batch, self.image_key) x_reconstructed, codebook_loss, loss_break = self(x) # generator aeloss, log_dict_ae = self.loss(codebook_loss, loss_break, x, x_reconstructed, 0, self.global_step, last_layer=self.get_last_layer(), split="train") self.manual_backward(aeloss / self.training_args.grad_accum_steps, inputs=self.generator_params) # https://discuss.pytorch.org/t/how-to-implement-gradient-accumulation-for-gan/112751/4 # discriminator discloss, log_dict_disc = self.loss(codebook_loss, loss_break, x, x_reconstructed, 1, self.global_step, last_layer=self.get_last_layer(), split="train") # x_reconstructed gets detached, `codebook_loss` and `loss_break` unused self.manual_backward(discloss / self.training_args.grad_accum_steps) # TODO: clip grads? if (batch_idx + 1) % self.training_args.grad_accum_steps == 0: # might not update at end of epoch? opt_gen, opt_disc = self.optimizers() scheduler_gen, scheduler_disc = self.lr_schedulers() #################### # fix global step bug # refer to https://github.com/Lightning-AI/pytorch-lightning/issues/17958 opt_disc._on_before_step = lambda: self.trainer.profiler.start("optimizer_step") opt_disc._on_after_step = lambda: self.trainer.profiler.stop("optimizer_step") # opt_gen._on_before_step = lambda: self.trainer.profiler.start("optimizer_step") # opt_gen._on_after_step = lambda: self.trainer.profiler.stop("optimizer_step") #################### opt_gen.step() scheduler_gen.step() opt_gen.zero_grad() opt_disc.step() scheduler_disc.step() opt_disc.zero_grad() self.log_dict(log_dict_disc, prog_bar=False, logger=True, on_step=True, on_epoch=True) self.log_dict(log_dict_ae, prog_bar=False, logger=True, on_step=True, on_epoch=True) def on_train_batch_end(self, *args, **kwargs): if self.use_ema: self.model_ema(self) def validation_step(self, batch, batch_idx): if self.use_ema: with self.ema_scope(): log_dict_ema = self._validation_step(batch, batch_idx, suffix="_ema") else: log_dict = self._validation_step(batch, batch_idx) def _validation_step(self, batch, batch_idx, suffix=""): x = self.get_input(batch, self.image_key) quant, eloss, indices, loss_break = self.encode(x) x_rec = self.decode(quant).clamp(-1, 1) aeloss, log_dict_ae = self.loss(eloss, loss_break, x, x_rec, 0, self.global_step, last_layer=self.get_last_layer(), split="val" + suffix) discloss, log_dict_disc = self.loss(eloss, loss_break, x, x_rec, 1, self.global_step, last_layer=self.get_last_layer(), split="val" + suffix) self.log_dict(log_dict_ae, prog_bar=False, logger=True, on_step=True, on_epoch=True) self.log_dict(log_dict_disc, prog_bar=False, logger=True, on_step=True, on_epoch=True) return self.log_dict def configure_optimizers(self): lr = self.training_args.learning_rate adam_betas = (self.training_args.adam_beta_1, self.training_args.adam_beta_2) opt_gen = torch.optim.Adam(self.generator_params, lr=lr, betas=adam_betas) opt_disc = torch.optim.Adam(self.loss.discriminator.parameters(), lr=lr, betas=adam_betas) # steps_per_epoch = len(self.trainer.datamodule._train_dataloader()) // self.trainer.world_size steps_per_epoch = len(self.trainer.fit_loop._data_source.instance) // self.trainer.world_size // self.training_args.grad_accum_steps if self.trainer.is_global_zero: print(f"{steps_per_epoch=}") warmup_steps = steps_per_epoch * self.training_args.warmup_epochs training_steps = steps_per_epoch * self.trainer.max_epochs if self.training_args.scheduler_type == "None": return ({"optimizer": opt_gen}, {"optimizer": opt_disc}) if self.training_args.scheduler_type == "linear-warmup": scheduler_ae = torch.optim.lr_scheduler.LambdaLR(opt_gen, Scheduler_LinearWarmup(warmup_steps)) scheduler_disc = torch.optim.lr_scheduler.LambdaLR(opt_disc, Scheduler_LinearWarmup(warmup_steps)) elif self.training_args.scheduler_type == "linear-warmup_cosine-decay": multipler_min = self.training_args.min_learning_rate / self.training_args.learning_rate scheduler_ae = torch.optim.lr_scheduler.LambdaLR(opt_gen, Scheduler_LinearWarmup_CosineDecay( warmup_steps=warmup_steps, max_steps=training_steps, multipler_min=multipler_min)) scheduler_disc = torch.optim.lr_scheduler.LambdaLR(opt_disc, Scheduler_LinearWarmup_CosineDecay( warmup_steps=warmup_steps, max_steps=training_steps, multipler_min=multipler_min)) else: raise NotImplementedError() return {"optimizer": opt_gen, "lr_scheduler": scheduler_ae}, {"optimizer": opt_disc, "lr_scheduler": scheduler_disc} def get_last_layer(self): return self.decoder.conv_out.weight def log_images(self, batch, **kwargs): log = dict() x = self.get_input(batch, self.image_key) x = x.to(self.device) xrec, _ = self(x) if x.shape[1] > 3: # colorize with random projection assert xrec.shape[1] > 3 x = self.to_rgb(x) xrec = self.to_rgb(xrec) log["inputs"] = x log["reconstructions"] = xrec return log def to_rgb(self, x): assert self.image_key == "segmentation" if not hasattr(self, "colorize"): self.register_buffer("colorize", torch.randn(3, x.shape[1], 1, 1).to(x)) x = F.conv2d(x, weight=self.colorize) x = 2. * (x - x.min()) / (x.max() - x.min()) - 1. return x