49 lines
1.5 KiB
Python
49 lines
1.5 KiB
Python
import ignite.distributed as idist
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from omegaconf import OmegaConf
|
|
|
|
from loss.I2I.perceptual_loss import PerceptualLoss
|
|
from loss.gan import GANLoss
|
|
|
|
|
|
def gan_loss(config):
|
|
gan_loss_cfg = OmegaConf.to_container(config)
|
|
gan_loss_cfg.pop("weight")
|
|
return GANLoss(**gan_loss_cfg).to(idist.device())
|
|
|
|
|
|
def perceptual_loss(config):
|
|
perceptual_loss_cfg = OmegaConf.to_container(config)
|
|
perceptual_loss_cfg.pop("weight")
|
|
return PerceptualLoss(**perceptual_loss_cfg).to(idist.device())
|
|
|
|
|
|
def pixel_loss(level):
|
|
return nn.L1Loss() if level == 1 else nn.MSELoss()
|
|
|
|
|
|
def mse_loss(x, target_flag):
|
|
return F.mse_loss(x, torch.ones_like(x) if target_flag else torch.zeros_like(x))
|
|
|
|
|
|
def bce_loss(x, target_flag):
|
|
return F.binary_cross_entropy_with_logits(x, torch.ones_like(x) if target_flag else torch.zeros_like(x))
|
|
|
|
|
|
def feature_match_loss(level, weight_policy):
|
|
compare_loss = pixel_loss(level)
|
|
assert weight_policy in ["same", "exponential_decline"]
|
|
|
|
def fm_loss(generated_features, target_features):
|
|
num_scale = len(generated_features)
|
|
loss = torch.zeros(1, device=idist.device())
|
|
for s_i in range(num_scale):
|
|
for i in range(len(generated_features[s_i]) - 1):
|
|
weight = 1 if weight_policy == "same" else 2 ** i
|
|
loss += weight * compare_loss(generated_features[s_i][i], target_features[s_i][i].detach()) / num_scale
|
|
return loss
|
|
|
|
return fm_loss
|