budui/raycv
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

v2

Browse Source
This commit is contained in:
budui
2020-10-22 22:42:01 +08:00
parent 0019d4034c
commit 376f5caeb7
11 changed files with 140 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
engine/CycleGAN.py
View File

@@ -1,11 +1,13 @@
from itertools import chain
import ignite.distributed as idist
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 pixel_loss, gan_loss
from engine.util.loss import pixel_loss, gan_loss, feature_match_loss
from loss.I2I.edge_loss import EdgeLoss
from loss.I2I.minimal_geometry_distortion_constraint_loss import MGCLoss
from model.weight_init import generation_init_weights
@@ -17,7 +19,10 @@ class CycleGANEngineKernel(EngineKernel):
self.gan_loss = gan_loss(config.loss.gan)
self.cycle_loss = LossContainer(config.loss.cycle.weight, pixel_loss(config.loss.cycle.level))
self.id_loss = LossContainer(config.loss.id.weight, pixel_loss(config.loss.id.level))
self.mgc_loss = LossContainer(config.loss.mgc.weight, MGCLoss())
self.mgc_loss = LossContainer(config.loss.mgc.weight, MGCLoss("opposite"))
self.fm_loss = LossContainer(config.loss.fm.weight, feature_match_loss(1, "same"))
self.edge_loss = LossContainer(config.loss.edge.weight, EdgeLoss(
"HED", hed_pretrained_model_path=config.loss.edge.hed_pretrained_model_path).to(idist.device()))
self.image_buffers = {k: GANImageBuffer(config.data.train.buffer_size or 50) for k in
self.discriminators.keys()}
@@ -64,8 +69,12 @@ class CycleGANEngineKernel(EngineKernel):
loss[f"cycle_{ph}"] = self.cycle_loss(generated["a2b2a" if ph == "a" else "b2a2b"], batch[ph])
loss[f"id_{ph}"] = self.id_loss(generated[f"{ph}2{ph}"], batch[ph])
loss[f"mgc_{ph}"] = self.mgc_loss(generated["a2b" if ph == "a" else "b2a"], batch[ph])
loss[f"gan_{ph}"] = self.config.loss.gan.weight * self.gan_loss(
self.discriminators[ph](generated["a2b" if ph == "b" else "b2a"]), True)
prediction_fake = self.discriminators[ph](generated["a2b" if ph == "b" else "b2a"])
loss[f"gan_{ph}"] = self.config.loss.gan.weight * self.gan_loss(prediction_fake, True)
if self.fm_loss.weight > 0:
prediction_real = self.discriminators[ph](batch[ph])
loss[f"feature_match_{ph}"] = self.fm_loss(prediction_fake, prediction_real)
loss[f"edge_{ph}"] = self.edge_loss(generated["a2b" if ph == "a" else "b2a"], batch[ph], gt_is_edge=False)
return loss
def criterion_discriminators(self, batch, generated) -> dict: