from itertools import chain import ignite.distributed as idist import torch import torch.nn as nn from omegaconf import OmegaConf from engine.base.i2i import EngineKernel, run_kernel from engine.util.build import build_model from loss.I2I.edge_loss import EdgeLoss from loss.I2I.perceptual_loss import PerceptualLoss from loss.gan import GANLoss from model.weight_init import generation_init_weights class TAFGEngineKernel(EngineKernel): def __init__(self, config): super().__init__(config) perceptual_loss_cfg = OmegaConf.to_container(config.loss.perceptual) perceptual_loss_cfg.pop("weight") self.perceptual_loss = PerceptualLoss(**perceptual_loss_cfg).to(idist.device()) gan_loss_cfg = OmegaConf.to_container(config.loss.gan) gan_loss_cfg.pop("weight") self.gan_loss = GANLoss(**gan_loss_cfg).to(idist.device()) self.recon_loss = nn.L1Loss() if config.loss.recon.level == 1 else nn.MSELoss() self.style_recon_loss = nn.L1Loss() if config.loss.style_recon.level == 1 else nn.MSELoss() self.content_recon_loss = nn.L1Loss() if config.loss.content_recon.level == 1 else nn.MSELoss() self.cycle_loss = nn.L1Loss() if config.loss.cycle.level == 1 else nn.MSELoss() self.edge_loss = EdgeLoss("HED", hed_pretrained_model_path=config.loss.edge.hed_pretrained_model_path).to( idist.device()) def _process_batch(self, batch, inference=False): # batch["b"] = batch["b"] if inference else batch["b"][0].expand(batch["a"].size()) return batch def build_models(self) -> (dict, dict): generators = dict( main=build_model(self.config.model.generator) ) discriminators = dict( a=build_model(self.config.model.discriminator), b=build_model(self.config.model.discriminator) ) self.logger.debug(discriminators["a"]) self.logger.debug(generators["main"]) for m in chain(generators.values(), discriminators.values()): generation_init_weights(m) return generators, discriminators def setup_after_g(self): for discriminator in self.discriminators.values(): discriminator.requires_grad_(True) def setup_before_g(self): for discriminator in self.discriminators.values(): discriminator.requires_grad_(False) def forward(self, batch, inference=False) -> dict: generator = self.generators["main"] batch = self._process_batch(batch, inference) styles = dict() contents = dict() images = dict() with torch.set_grad_enabled(not inference): for ph in "ab": contents[ph], styles[ph] = generator.encode(batch[ph]["edge"], batch[ph]["img"], ph, ph) for ph in ("a2b", "b2a"): images[f"fake_{ph[-1]}"] = generator.decode(contents[ph[0]], styles[ph[-1]], ph[-1]) contents["recon_a"], styles["recon_b"] = generator.encode( self.edge_loss.edge_extractor(images["fake_b"]), images["fake_b"], "b", "b") images["a2a"] = generator.decode(contents["a"], styles["a"], "a") images["b2b"] = generator.decode(contents["b"], styles["recon_b"], "b") images["cycle_a"] = generator.decode(contents["recon_a"], styles["a"], "a") return dict(styles=styles, contents=contents, images=images) def criterion_generators(self, batch, generated) -> dict: batch = self._process_batch(batch) loss = dict() for ph in "ab": loss[f"recon_image_{ph}"] = self.config.loss.recon.weight * self.recon_loss( generated["images"][f"{ph}2{ph}"], batch[ph]["img"]) pred_fake = self.discriminators[ph](generated["images"][f"fake_{ph}"]) loss[f"gan_{ph}"] = 0 for sub_pred_fake in pred_fake: # last output is actual prediction loss[f"gan_{ph}"] += self.gan_loss(sub_pred_fake[-1], True) * self.config.loss.gan.weight loss[f"recon_content_a"] = self.config.loss.content_recon.weight * self.content_recon_loss( generated["contents"]["a"], generated["contents"]["recon_a"] ) loss[f"recon_style_b"] = self.config.loss.style_recon.weight * self.style_recon_loss( generated["styles"]["b"], generated["styles"]["recon_b"] ) for ph in ("a2b", "b2a"): if self.config.loss.perceptual.weight > 0: loss[f"perceptual_{ph}"] = self.config.loss.perceptual.weight * self.perceptual_loss( batch[ph[0]]["img"], generated["images"][f"fake_{ph[-1]}"] ) if self.config.loss.edge.weight > 0: loss[f"edge_a"] = self.config.loss.edge.weight * self.edge_loss( generated["images"]["fake_b"], batch["a"]["edge"][:, 0:1, :, :] ) loss[f"edge_b"] = self.config.loss.edge.weight * self.edge_loss( generated["images"]["fake_a"], batch["b"]["edge"] ) if self.config.loss.cycle.weight > 0: loss[f"cycle_a"] = self.config.loss.cycle.weight * self.cycle_loss( batch["a"]["img"], generated["images"]["cycle_a"] ) return loss def criterion_discriminators(self, batch, generated) -> dict: loss = dict() # batch = self._process_batch(batch) for phase in self.discriminators.keys(): pred_real = self.discriminators[phase](batch[phase]["img"]) pred_fake = self.discriminators[phase](generated["images"][f"fake_{phase}"].detach()) loss[f"gan_{phase}"] = 0 for i in range(len(pred_fake)): loss[f"gan_{phase}"] += (self.gan_loss(pred_fake[i][-1], False, is_discriminator=True) + self.gan_loss(pred_real[i][-1], True, is_discriminator=True)) / 2 return loss def intermediate_images(self, batch, generated) -> dict: """ returned dict must be like: {"a": [img1, img2, ...], "b": [img3, img4, ...]} :param batch: :param generated: dict of images :return: dict like: {"a": [img1, img2, ...], "b": [img3, img4, ...]} """ batch = self._process_batch(batch) return dict( a=[batch["a"]["edge"][:, 0:1, :, :].expand(-1, 3, -1, -1).detach(), batch["a"]["img"].detach(), generated["images"]["a2a"].detach(), generated["images"]["fake_b"].detach(), generated["images"]["cycle_a"].detach(), ], b=[batch["b"]["edge"].expand(-1, 3, -1, -1).detach(), batch["b"]["img"].detach(), generated["images"]["b2b"].detach(), generated["images"]["fake_a"].detach()] ) def change_engine(self, config, trainer): pass # @trainer.on(Events.ITERATION_STARTED(once=int(config.max_iteration / 3))) # def change_config(engine): # with read_write(config): # config.loss.perceptual.weight = 5 def run(task, config, _): kernel = TAFGEngineKernel(config) run_kernel(task, config, kernel)