def __init__(self,
recursions=1,
stride=1,
kernel_size=5,
use_perceptual_loss=True,
wgan=False,
w_col=1,
w_tex=0.001,
w_per=0.1,
gaussian=False,
lpips_rot_flip=False,
**kwargs):
super(GeneratorLoss, self).__init__()
self.pixel_loss = nn.L1Loss()
self.color_filter = FilterLow(recursions=recursions,
stride=stride,
kernel_size=kernel_size,
padding=False,
gaussian=gaussian)
if torch.cuda.is_available():
self.pixel_loss = self.pixel_loss.cuda()
self.color_filter = self.color_filter.cuda()
self.perceptual_loss = PerceptualLoss(rotations=lpips_rot_flip,
flips=lpips_rot_flip)
self.use_perceptual_loss = use_perceptual_loss
self.wasserstein = wgan
self.w_col = w_col
self.w_tex = w_tex
self.w_per = w_per
self.last_tex_loss = 0
self.last_per_loss = 0
self.last_col_loss = 0
self.last_mean_loss = 0
def __init__(self,
recursions=1,
stride=1,
kernel_size=5,
use_perceptual_loss=True,
wgan=False,
w_col=1,
w_tex=0.001,
w_per=0.1,
gaussian=False,
lpips_rot_flip=False,
**kwargs):
super(GeneratorLoss, self).__init__()
self.pixel_loss = nn.L1Loss()
self.per_type = kwargs['per_type']
if kwargs['filter'].lower() == 'gau':
self.color_filter = FilterLow(recursions=recursions,
stride=stride,
kernel_size=kernel_size,
padding=False,
gaussian=True)
elif kwargs['filter'].lower() == 'avg_pool':
self.color_filter = FilterLow(recursions=recursions,
stride=stride,
kernel_size=kernel_size,
padding=False,
gaussian=False)
elif kwargs['filter'].lower() == 'wavelet':
self.color_filter = self.filter_wavelet_LL
else:
raise NotImplementedError(
'Frequency Separation type [{:s}] not recognized'.format(
kwargs['filter']))
if torch.cuda.is_available():
self.pixel_loss = self.pixel_loss.cuda()
if isinstance(self.color_filter, FilterLow):
self.color_filter = self.color_filter.cuda()
if self.per_type == 'LPIPS':
self.perceptual_loss = PerceptualLoss(rotations=lpips_rot_flip,
flips=lpips_rot_flip)
elif self.per_type == 'VGG':
self.perceptual_loss = PerceptualLossVGG16()
else:
raise NotImplemented('{} is not recognized'.format(self.per_type))
self.use_perceptual_loss = use_perceptual_loss
self.wasserstein = wgan
self.w_col = w_col
self.w_tex = w_tex
self.w_per = w_per
self.last_tex_loss = 0
self.last_per_loss = 0
self.last_col_loss = 0
self.gaussian = gaussian
self.last_mean_loss = 0
class GeneratorLoss(nn.Module):
def __init__(self,
recursions=1,
stride=1,
kernel_size=5,
use_perceptual_loss=True,
wgan=False,
w_col=1,
w_tex=0.001,
w_per=0.1,
gaussian=False,
lpips_rot_flip=False,
**kwargs):
super(GeneratorLoss, self).__init__()
self.pixel_loss = nn.L1Loss()
self.color_filter = FilterLow(recursions=recursions,
stride=stride,
kernel_size=kernel_size,
padding=False,
gaussian=gaussian)
if torch.cuda.is_available():
self.pixel_loss = self.pixel_loss.cuda()
self.color_filter = self.color_filter.cuda()
self.perceptual_loss = PerceptualLoss(rotations=lpips_rot_flip,
flips=lpips_rot_flip)
self.use_perceptual_loss = use_perceptual_loss
self.wasserstein = wgan
self.w_col = w_col
self.w_tex = w_tex
self.w_per = w_per
self.last_tex_loss = 0
self.last_per_loss = 0
self.last_col_loss = 0
self.last_mean_loss = 0
def forward(self, tex_labels, out_images, target_images):
# Adversarial Texture Loss
self.last_tex_loss = generator_loss(tex_labels,
wasserstein=self.wasserstein)
# Perception Loss
self.last_per_loss = self.perceptual_loss(out_images, target_images)
# Color Loss
self.last_col_loss = self.color_loss(out_images, target_images)
loss = self.w_col * self.last_col_loss + self.w_tex * self.last_tex_loss
if self.use_perceptual_loss:
loss += self.w_per * self.last_per_loss
return loss
def color_loss(self, x, y):
return self.pixel_loss(self.color_filter(x), self.color_filter(y))
def rgb_loss(self, x, y):
return self.pixel_loss(x.mean(3).mean(2), y.mean(3).mean(2))
def mean_loss(self, x, y):
return self.pixel_loss(
x.view(x.size(0), -1).mean(1),
y.view(y.size(0), -1).mean(1))
class GeneratorLoss(nn.Module):
def __init__(self,
recursions=1,
stride=1,
kernel_size=5,
use_perceptual_loss=True,
wgan=False,
w_col=1,
w_tex=0.001,
w_per=0.1,
gaussian=False,
lpips_rot_flip=False,
**kwargs):
super(GeneratorLoss, self).__init__()
self.pixel_loss = nn.L1Loss()
self.per_type = kwargs['per_type']
if kwargs['filter'].lower() == 'gau':
self.color_filter = FilterLow(recursions=recursions,
stride=stride,
kernel_size=kernel_size,
padding=False,
gaussian=True)
elif kwargs['filter'].lower() == 'avg_pool':
self.color_filter = FilterLow(recursions=recursions,
stride=stride,
kernel_size=kernel_size,
padding=False,
gaussian=False)
elif kwargs['filter'].lower() == 'wavelet':
self.color_filter = self.filter_wavelet_LL
else:
raise NotImplementedError(
'Frequency Separation type [{:s}] not recognized'.format(
kwargs['filter']))
if torch.cuda.is_available():
self.pixel_loss = self.pixel_loss.cuda()
if isinstance(self.color_filter, FilterLow):
self.color_filter = self.color_filter.cuda()
if self.per_type == 'LPIPS':
self.perceptual_loss = PerceptualLoss(rotations=lpips_rot_flip,
flips=lpips_rot_flip)
elif self.per_type == 'VGG':
self.perceptual_loss = PerceptualLossVGG16()
else:
raise NotImplemented('{} is not recognized'.format(self.per_type))
self.use_perceptual_loss = use_perceptual_loss
self.wasserstein = wgan
self.w_col = w_col
self.w_tex = w_tex
self.w_per = w_per
self.last_tex_loss = 0
self.last_per_loss = 0
self.last_col_loss = 0
self.gaussian = gaussian
self.last_mean_loss = 0
def forward(self, tex_labels, out_images, target_images):
# Adversarial Texture Loss
self.last_tex_loss = generator_loss(tex_labels,
wasserstein=self.wasserstein)
# Perception Loss
self.last_per_loss = self.perceptual_loss(out_images, target_images)
# Color Loss
self.last_col_loss = self.color_loss(out_images, target_images)
loss = self.w_col * self.last_col_loss + self.w_tex * self.last_tex_loss
if self.use_perceptual_loss:
loss += self.w_per * self.last_per_loss
return loss
def color_loss(self, x, y):
return self.pixel_loss(self.color_filter(x), self.color_filter(y))
def rgb_loss(self, x, y):
return self.pixel_loss(x.mean(3).mean(2), y.mean(3).mean(2))
def mean_loss(self, x, y):
return self.pixel_loss(
x.view(x.size(0), -1).mean(1),
y.view(y.size(0), -1).mean(1))
def filter_wavelet_LL(self, x, norm=True):
DWT2 = DWTForward(J=1, wave='haar', mode='reflect').cuda()
LL, Hc = DWT2(x)
return LL * 0.5 if norm else LL