# Copyright (C) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# This work is made available under the Nvidia Source Code License-NC.
# To view a copy of this license, check out LICENSE.md
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from imaginaire.utils.distributed import master_only_print as print
@torch.jit.script
def fuse_math_min_mean_pos(x):
r"""Fuse operation min mean for hinge loss computation of positive
samples"""
minval = torch.min(x - 1, x * 0)
loss = -torch.mean(minval)
return loss
@torch.jit.script
def fuse_math_min_mean_neg(x):
r"""Fuse operation min mean for hinge loss computation of negative
samples"""
minval = torch.min(-x - 1, x * 0)
loss = -torch.mean(minval)
return loss
[docs]class GANLoss(nn.Module):
r"""GAN loss constructor.
Args:
gan_mode (str): Type of GAN loss. ``'hinge'``, ``'least_square'``,
``'non_saturated'``, ``'wasserstein'``.
target_real_label (float): The desired output label for real images.
target_fake_label (float): The desired output label for fake images.
decay_k (float): The decay factor per epoch for top-k training.
min_k (float): The minimum percentage of samples to select.
separate_topk (bool): If ``True``, selects top-k for each sample
separately, otherwise selects top-k among all samples.
"""
def __init__(self, gan_mode, target_real_label=1.0, target_fake_label=0.0,
decay_k=1., min_k=1., separate_topk=False):
super(GANLoss, self).__init__()
self.real_label = target_real_label
self.fake_label = target_fake_label
self.real_label_tensor = None
self.fake_label_tensor = None
self.gan_mode = gan_mode
self.decay_k = decay_k
self.min_k = min_k
self.separate_topk = separate_topk
self.register_buffer('k', torch.tensor(1.0))
print('GAN mode: %s' % gan_mode)
[docs] def forward(self, dis_output, t_real, dis_update=True, reduce=True):
r"""GAN loss computation.
Args:
dis_output (tensor or list of tensors): Discriminator outputs.
t_real (bool): If ``True``, uses the real label as target, otherwise uses the fake label as target.
dis_update (bool): If ``True``, the loss will be used to update the discriminator, otherwise the generator.
reduce (bool): If ``True``, when a list of discriminator outputs are provided, it will return the average
of all losses, otherwise it will return a list of losses.
Returns:
loss (tensor): Loss value.
"""
if isinstance(dis_output, list):
# For multi-scale discriminators.
# In this implementation, the loss is first averaged for each scale
# (batch size and number of locations) then averaged across scales,
# so that the gradient is not dominated by the discriminator that
# has the most output values (highest resolution).
losses = []
for dis_output_i in dis_output:
assert isinstance(dis_output_i, torch.Tensor)
losses.append(self.loss(dis_output_i, t_real, dis_update))
if reduce:
return torch.mean(torch.stack(losses))
else:
return losses
else:
return self.loss(dis_output, t_real, dis_update)
[docs] def loss(self, dis_output, t_real, dis_update=True):
r"""GAN loss computation.
Args:
dis_output (tensor): Discriminator outputs.
t_real (bool): If ``True``, uses the real label as target, otherwise
uses the fake label as target.
dis_update (bool): Updating the discriminator or the generator.
Returns:
loss (tensor): Loss value.
"""
if not dis_update:
assert t_real, \
"The target should be real when updating the generator."
if not dis_update and self.k < 1:
r"""
Use top-k training:
"Top-k Training of GANs: Improving GAN Performance by Throwing
Away Bad Samples"
Here, each sample may have multiple discriminator output values
(patch discriminator). We could either select top-k for each sample
separately (when ``self.separate_topk=True``), or collect values
from all samples and then select top-k (default, when
``self.separate_topk=False``).
"""
if self.separate_topk:
dis_output = dis_output.view(dis_output.size(0), -1)
else:
dis_output = dis_output.view(-1)
k = math.ceil(self.k * dis_output.size(-1))
dis_output, _ = torch.topk(dis_output, k)
if self.gan_mode == 'non_saturated':
target_tensor = self.get_target_tensor(dis_output, t_real)
loss = F.binary_cross_entropy_with_logits(dis_output,
target_tensor)
elif self.gan_mode == 'least_square':
target_tensor = self.get_target_tensor(dis_output, t_real)
loss = 0.5 * F.mse_loss(dis_output, target_tensor)
elif self.gan_mode == 'hinge':
if dis_update:
if t_real:
loss = fuse_math_min_mean_pos(dis_output)
else:
loss = fuse_math_min_mean_neg(dis_output)
else:
loss = -torch.mean(dis_output)
elif self.gan_mode == 'wasserstein':
if t_real:
loss = -torch.mean(dis_output)
else:
loss = torch.mean(dis_output)
elif self.gan_mode == 'softplus':
target_tensor = self.get_target_tensor(dis_output, t_real)
loss = F.binary_cross_entropy_with_logits(dis_output,
target_tensor)
else:
raise ValueError('Unexpected gan_mode {}'.format(self.gan_mode))
return loss
[docs] def get_target_tensor(self, dis_output, t_real):
r"""Return the target vector for the binary cross entropy loss
computation.
Args:
dis_output (tensor): Discriminator outputs.
t_real (bool): If ``True``, uses the real label as target, otherwise
uses the fake label as target.
Returns:
target (tensor): Target tensor vector.
"""
if t_real:
if self.real_label_tensor is None:
self.real_label_tensor = dis_output.new_tensor(self.real_label)
return self.real_label_tensor.expand_as(dis_output)
else:
if self.fake_label_tensor is None:
self.fake_label_tensor = dis_output.new_tensor(self.fake_label)
return self.fake_label_tensor.expand_as(dis_output)
[docs] def topk_anneal(self):
r"""Anneal k after each epoch."""
if self.decay_k < 1:
# noinspection PyAttributeOutsideInit
self.k.fill_(max(self.decay_k * self.k, self.min_k))
print("Top-k training: update k to {}.".format(self.k))