87 lines
3.4 KiB
Python
87 lines
3.4 KiB
Python
from itertools import chain
|
|
|
|
import torch
|
|
|
|
from engine.base.i2i import EngineKernel, run_kernel
|
|
from engine.util.build import build_model
|
|
from engine.util.container import GANImageBuffer, LossContainer
|
|
from engine.util.loss import gan_loss, feature_match_loss, perceptual_loss
|
|
from loss.I2I.minimal_geometry_distortion_constraint_loss import MGCLoss
|
|
from model.weight_init import generation_init_weights
|
|
|
|
|
|
class GauGANEngineKernel(EngineKernel):
|
|
def __init__(self, config):
|
|
super().__init__(config)
|
|
|
|
self.gan_loss = gan_loss(config.loss.gan)
|
|
self.mgc_loss = LossContainer(config.loss.mgc.weight, MGCLoss("opposite"))
|
|
self.fm_loss = LossContainer(config.loss.fm.weight, feature_match_loss(1, "exponential_decline"))
|
|
self.perceptual_loss = LossContainer(config.loss.perceptual.weight, perceptual_loss(config.loss.perceptual))
|
|
|
|
self.image_buffers = {k: GANImageBuffer(config.data.train.buffer_size or 50) for k in
|
|
self.discriminators.keys()}
|
|
|
|
def build_models(self) -> (dict, dict):
|
|
generators = dict(
|
|
main=build_model(self.config.model.generator)
|
|
)
|
|
discriminators = dict(
|
|
b=build_model(self.config.model.discriminator)
|
|
)
|
|
self.logger.debug(discriminators["b"])
|
|
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:
|
|
images = dict()
|
|
with torch.set_grad_enabled(not inference):
|
|
images["a2b"] = self.generators["main"](batch["a"])
|
|
return images
|
|
|
|
def criterion_generators(self, batch, generated) -> dict:
|
|
loss = dict()
|
|
prediction_fake = self.discriminators["b"](generated["a2b"])
|
|
loss["gan"] = self.config.loss.gan.weight * self.gan_loss(prediction_fake, True)
|
|
loss["mgc"] = self.mgc_loss(generated["a2b"], batch["a"])
|
|
loss["perceptual"] = self.perceptual_loss(generated["a2b"], batch["a"])
|
|
if self.fm_loss.weight > 0:
|
|
prediction_real = self.discriminators["b"](batch["b"])
|
|
loss["feature_match"] = self.fm_loss(prediction_fake, prediction_real)
|
|
return loss
|
|
|
|
def criterion_discriminators(self, batch, generated) -> dict:
|
|
loss = dict()
|
|
generated_image = self.image_buffers["b"].query(generated["a2b"].detach())
|
|
loss["b"] = (self.gan_loss(self.discriminators["b"](generated_image), False, is_discriminator=True) +
|
|
self.gan_loss(self.discriminators["b"](batch["b"]), 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, ...]}
|
|
"""
|
|
return dict(
|
|
a=[batch["a"].detach(), generated["a2b"].detach()],
|
|
)
|
|
|
|
|
|
def run(task, config, _):
|
|
kernel = GauGANEngineKernel(config)
|
|
run_kernel(task, config, kernel)
|