mmagic.models.losses
¶
Package Contents¶
Classes¶
Returns the loss of discriminator with the specified type for 

Clip loss. In styleclip, this loss is used to optimize the latent code 

Charbonnier composition loss. 

L1 composition loss. 

MSE (L2) composition loss. 

Face similarity loss. Generally this loss is used to keep the id 

Feature loss of DICGAN, based on LightCNN. 

Disc shift loss. 

Define GAN loss. 

A Gaussian filter which blurs a given tensor with a twodimensional 

Gradient penalty loss for wgangp. 

Gradient loss. 

Clip loss. In styleclip, this loss is used to optimize the latent code 

Disc Shift Loss. 

Face similarity loss. Generally this loss is used to keep the id 

Define GAN loss. 

Generator Path Regularizer. 

Gradient Penalty for WGANGP. 

R1 gradient penalty for WGANGP. 

Perceptual loss with commonly used style loss. 

VGG network used in calculating perceptual loss. 

Transferal perceptual loss. 

Charbonnier loss (one variant of Robust L1Loss, a differentiable variant 

L1 (mean absolute error, MAE) loss. 

Masked TV loss. 

MSE (L2) loss. 

PSNR Loss in "HINet: Half Instance Normalization Network for Image 
Functions¶

Disc Shift loss. 

Generator Path Regularization. 

Calculate gradient penalty for wgangp. 

Calculate R1 gradient penalty for WGANGP. 

Apply elementwise weight and reduce loss. 

Reduce loss as specified. 

L2 total variation loss, as in Mahendran et al. 
 class mmagic.models.losses.AdvLoss[source]¶
Returns the loss of discriminator with the specified type for DeblurGanv2.
 Parameters
loss_type (Str) – One of value in [wgangp,lsgan,gan,ragan,raganls].
 class mmagic.models.losses.CLIPLoss(loss_weight: float = 1.0, data_info: Optional[dict] = None, clip_model: dict = dict(), loss_name: str = 'loss_clip')[source]¶
Bases:
torch.nn.Module
Clip loss. In styleclip, this loss is used to optimize the latent code to generate image that match the text.
In this loss, we may need to provide
image
,text
. Thus, an example of thedata_info
is:1data_info = dict( 2 image='fake_imgs', 3 text='descriptions')
Then, the module will automatically construct this mapping from the input data dictionary.
 Parameters
loss_weight (float, optional) – Weight of this loss item. Defaults to
1.
.data_info (dict, optional) – Dictionary contains the mapping between loss input args and data dictionary. If
None
, this module will directly pass the input data to the loss function. Defaults to None.clip_model (dict, optional) – Kwargs for clip loss model. Defaults to dict().
loss_name (str, optional) – Name of the loss item. If you want this loss item to be included into the backward graph, loss_ must be the prefix of the name. Defaults to ‘loss_clip’.
 forward(image: torch.Tensor, text: torch.Tensor) torch.Tensor [source]¶
Forward function.
If
self.data_info
is notNone
, a dictionary containing all of the data and necessary modules should be passed into this function. If this dictionary is given as a nonkeyword argument, it should be offered as the first argument. If you are using keyword argument, please name it as outputs_dict.If
self.data_info
isNone
, the input argument or keyword argument will be directly passed to loss function,third_party_net_loss
.
 class mmagic.models.losses.CharbonnierCompLoss(loss_weight: float = 1.0, reduction: str = 'mean', sample_wise: bool = False, eps: bool = 1e12)[source]¶
Bases:
torch.nn.Module
Charbonnier composition loss.
 Parameters
loss_weight (float) – Loss weight for L1 loss. Default: 1.0.
reduction (str) – Specifies the reduction to apply to the output. Supported choices are ‘none’  ‘mean’  ‘sum’. Default: ‘mean’.
sample_wise (bool) – Whether calculate the loss samplewise. This argument only takes effect when reduction is ‘mean’ and weight (argument of forward()) is not None. It will first reduces loss with ‘mean’ persample, and then it means over all the samples. Default: False.
eps (float) – A value used to control the curvature near zero. Default: 1e12.
 forward(pred_alpha: torch.Tensor, fg: torch.Tensor, bg: torch.Tensor, ori_merged: torch.Tensor, weight: Optional[torch.Tensor] = None, **kwargs) torch.Tensor [source]¶
 Parameters
pred_alpha (Tensor) – of shape (N, 1, H, W). Predicted alpha matte.
fg (Tensor) – of shape (N, 3, H, W). Tensor of foreground object.
bg (Tensor) – of shape (N, 3, H, W). Tensor of background object.
ori_merged (Tensor) – of shape (N, 3, H, W). Tensor of origin merged image before normalized by ImageNet mean and std.
weight (Tensor, optional) – of shape (N, 1, H, W). It is an indicating matrix: weight[trimap == 128] = 1. Default: None.
 class mmagic.models.losses.L1CompositionLoss(loss_weight: float = 1.0, reduction: str = 'mean', sample_wise: bool = False)[source]¶
Bases:
torch.nn.Module
L1 composition loss.
 Parameters
loss_weight (float) – Loss weight for L1 loss. Default: 1.0.
reduction (str) – Specifies the reduction to apply to the output. Supported choices are ‘none’  ‘mean’  ‘sum’. Default: ‘mean’.
sample_wise (bool) – Whether calculate the loss samplewise. This argument only takes effect when reduction is ‘mean’ and weight (argument of forward()) is not None. It will first reduces loss with ‘mean’ persample, and then it means over all the samples. Default: False.
 forward(pred_alpha: torch.Tensor, fg: torch.Tensor, bg: torch.Tensor, ori_merged: torch.Tensor, weight: Optional[torch.Tensor] = None, **kwargs) torch.Tensor [source]¶
 Parameters
pred_alpha (Tensor) – of shape (N, 1, H, W). Predicted alpha matte.
fg (Tensor) – of shape (N, 3, H, W). Tensor of foreground object.
bg (Tensor) – of shape (N, 3, H, W). Tensor of background object.
ori_merged (Tensor) – of shape (N, 3, H, W). Tensor of origin merged image before normalized by ImageNet mean and std.
weight (Tensor, optional) – of shape (N, 1, H, W). It is an indicating matrix: weight[trimap == 128] = 1. Default: None.
 class mmagic.models.losses.MSECompositionLoss(loss_weight: float = 1.0, reduction: str = 'mean', sample_wise: bool = False)[source]¶
Bases:
torch.nn.Module
MSE (L2) composition loss.
 Parameters
loss_weight (float) – Loss weight for MSE loss. Default: 1.0.
reduction (str) – Specifies the reduction to apply to the output. Supported choices are ‘none’  ‘mean’  ‘sum’. Default: ‘mean’.
sample_wise (bool) – Whether calculate the loss samplewise. This argument only takes effect when reduction is ‘mean’ and weight (argument of forward()) is not None. It will first reduces loss with ‘mean’ persample, and then it means over all the samples. Default: False.
 forward(pred_alpha: torch.Tensor, fg: torch.Tensor, bg: torch.Tensor, ori_merged: torch.Tensor, weight: Optional[torch.Tensor] = None, **kwargs) torch.Tensor [source]¶
 Parameters
pred_alpha (Tensor) – of shape (N, 1, H, W). Predicted alpha matte.
fg (Tensor) – of shape (N, 3, H, W). Tensor of foreground object.
bg (Tensor) – of shape (N, 3, H, W). Tensor of background object.
ori_merged (Tensor) – of shape (N, 3, H, W). Tensor of origin merged image before normalized by ImageNet mean and std.
weight (Tensor, optional) – of shape (N, 1, H, W). It is an indicating matrix: weight[trimap == 128] = 1. Default: None.
 class mmagic.models.losses.FaceIdLoss(loss_weight: float = 1.0, data_info: Optional[dict] = None, facenet: dict = dict(type='ArcFace', ir_se50_weights=None), loss_name: str = 'loss_id')[source]¶
Bases:
torch.nn.Module
Face similarity loss. Generally this loss is used to keep the id consistency of the input face image and output face image.
In this loss, we may need to provide
gt
,pred
andx
. Thus, an example of thedata_info
is:1data_info = dict( 2 gt='real_imgs', 3 pred='fake_imgs')
Then, the module will automatically construct this mapping from the input data dictionary.
 Parameters
loss_weight (float, optional) – Weight of this loss item. Defaults to
1.
.data_info (dict, optional) – Dictionary contains the mapping between loss input args and data dictionary. If
None
, this module will directly pass the input data to the loss function. Defaults to None.facenet (dict, optional) – Config dict for facenet. Defaults to dict(type=’ArcFace’, ir_se50_weights=None, device=’cuda’).
loss_name (str, optional) – Name of the loss item. If you want this loss item to be included into the backward graph, loss_ must be the prefix of the name. Defaults to ‘loss_id’.
 class mmagic.models.losses.LightCNNFeatureLoss(pretrained: str, loss_weight: float = 1.0, criterion: str = 'l1')[source]¶
Bases:
torch.nn.Module
Feature loss of DICGAN, based on LightCNN.
 Parameters
pretrained (str) – Path for pretrained weights.
loss_weight (float) – Loss weight. Default: 1.0.
criterion (str) – Criterion type. Options are ‘l1’ and ‘mse’. Default: ‘l1’.
 class mmagic.models.losses.DiscShiftLoss(loss_weight: float = 0.1)[source]¶
Bases:
torch.nn.Module
Disc shift loss.
 Parameters
loss_weight (float, optional) – Loss weight. Defaults to 1.0.
 class mmagic.models.losses.GANLoss(gan_type: str, real_label_val: float = 1.0, fake_label_val: float = 0.0, loss_weight: float = 1.0)[source]¶
Bases:
torch.nn.Module
Define GAN loss.
 Parameters
gan_type (str) – Support ‘vanilla’, ‘lsgan’, ‘wgan’, ‘hinge’, ‘l1’.
real_label_val (float) – The value for real label. Default: 1.0.
fake_label_val (float) – The value for fake label. Default: 0.0.
loss_weight (float) – Loss weight. Default: 1.0. Note that loss_weight is only for generators; and it is always 1.0 for discriminators.
 _wgan_loss(input: torch.Tensor, target: bool) torch.Tensor [source]¶
wgan loss.
 Parameters
input (Tensor) – Input tensor.
target (bool) – Target label.
 Returns
wgan loss.
 Return type
Tensor
 get_target_label(input: torch.Tensor, target_is_real: bool) Union[bool, torch.Tensor] [source]¶
Get target label.
 Parameters
input (Tensor) – Input tensor.
target_is_real (bool) – Whether the target is real or fake.
 Returns
 Target tensor. Return bool for wgan, otherwise,
return Tensor.
 Return type
(bool  Tensor)
 forward(input: torch.Tensor, target_is_real: bool, is_disc: bool = False, mask: Optional[torch.Tensor] = None) torch.Tensor [source]¶
 Parameters
input (Tensor) – The input for the loss module, i.e., the network prediction.
target_is_real (bool) – Whether the target is real or fake.
is_disc (bool) – Whether the loss for discriminators or not. Default: False.
mask (Tensor) – The mask tensor. Default: None.
 Returns
GAN loss value.
 Return type
Tensor
 class mmagic.models.losses.GaussianBlur(kernel_size: Tuple[int, int] = (71, 71), sigma: Tuple[float, float] = (10.0, 10.0))[source]¶
Bases:
torch.nn.Module
A Gaussian filter which blurs a given tensor with a twodimensional gaussian kernel by convolving it along each channel. Batch operation is supported.
This function is modified from kornia.filters.gaussian: <https://kornia.readthedocs.io/en/latest/_modules/kornia/filters/gaussian.html>.
 Parameters
kernel_size (tuple[int]) – The size of the kernel. Default: (71, 71).
sigma (tuple[float]) – The standard deviation of the kernel.
Default (10.0, 10.0) –
 Returns
The Gaussianblurred tensor.
 Return type
Tensor
 Shape:
input: Tensor with shape of (n, c, h, w)
output: Tensor with shape of (n, c, h, w)
 static compute_zero_padding(kernel_size: Tuple[int, int]) tuple [source]¶
Compute zero padding tuple.
 Parameters
kernel_size (tuple[int]) – The size of the kernel.
 Returns
Padding of height and weight.
 Return type
tuple
 get_2d_gaussian_kernel(kernel_size: Tuple[int, int], sigma: Tuple[float, float]) torch.Tensor [source]¶
Get the twodimensional Gaussian filter matrix coefficients.
 Parameters
kernel_size (tuple[int]) – Kernel filter size in the x and y direction. The kernel sizes should be odd and positive.
sigma (tuple[int]) – Gaussian standard deviation in the x and y direction.
 Returns
 A 2D torch tensor with gaussian filter
matrix coefficients.
 Return type
kernel_2d (Tensor)
 get_1d_gaussian_kernel(kernel_size: int, sigma: float) torch.Tensor [source]¶
Get the Gaussian filter coefficients in one dimension (x or y direction).
 Parameters
kernel_size (int) – Kernel filter size in x or y direction. Should be odd and positive.
sigma (float) – Gaussian standard deviation in x or y direction.
 Returns
 A 1D torch tensor with gaussian filter
coefficients in x or y direction.
 Return type
kernel_1d (Tensor)
 gaussian(kernel_size: int, sigma: float) torch.Tensor [source]¶
Gaussian function.
 Parameters
kernel_size (int) – Kernel filter size in x or y direction. Should be odd and positive.
sigma (float) – Gaussian standard deviation in x or y direction.
 Returns
 A 1D torch tensor with gaussian filter
coefficients in x or y direction.
 Return type
Tensor
 class mmagic.models.losses.GradientPenaltyLoss(loss_weight: float = 1.0)[source]¶
Bases:
torch.nn.Module
Gradient penalty loss for wgangp.
 Parameters
loss_weight (float) – Loss weight. Default: 1.0.
 forward(discriminator: torch.nn.Module, real_data: torch.Tensor, fake_data: torch.Tensor, mask: Optional[torch.Tensor] = None) torch.Tensor [source]¶
Forward function.
 Parameters
discriminator (nn.Module) – Network for the discriminator.
real_data (Tensor) – Real input data.
fake_data (Tensor) – Fake input data.
mask (Tensor) – Masks for inpainting. Default: None.
 Returns
Loss.
 Return type
Tensor
 mmagic.models.losses.disc_shift_loss(pred: torch.Tensor) torch.Tensor [source]¶
Disc Shift loss.
This loss is proposed in PGGAN as an auxiliary loss for discriminator.
 Parameters
pred (Tensor) – Input tensor.
 Returns
loss tensor.
 Return type
torch.Tensor
 mmagic.models.losses.gen_path_regularizer(generator: torch.nn.Module, num_batches: int, mean_path_length: torch.Tensor, pl_batch_shrink: int = 1, decay: float = 0.01, weight: float = 1.0, pl_batch_size: Optional[int] = None, sync_mean_buffer: bool = False, loss_scaler: Optional[torch.cuda.amp.grad_scaler.GradScaler] = None, use_apex_amp: bool = False) Tuple[torch.Tensor] [source]¶
Generator Path Regularization.
Path regularization is proposed in StyleGAN2, which can help the improve the continuity of the latent space. More details can be found in: Analyzing and Improving the Image Quality of StyleGAN, CVPR2020.
 Parameters
generator (nn.Module) – The generator module. Note that this loss requires that the generator contains
return_latents
interface, with which we can get the latent code of the current sample.num_batches (int) – The number of samples used in calculating this loss.
mean_path_length (Tensor) – The mean path length, calculated by moving average.
pl_batch_shrink (int, optional) – The factor of shrinking the batch size for saving GPU memory. Defaults to 1.
decay (float, optional) – Decay for moving average of mean path length. Defaults to 0.01.
weight (float, optional) – Weight of this loss item. Defaults to
1.
.pl_batch_size (int  None, optional) – The batch size in calculating generator path. Once this argument is set, the
num_batches
will be overridden with this argument and won’t be affected bypl_batch_shrink
. Defaults to None.sync_mean_buffer (bool, optional) – Whether to sync mean path length across all of GPUs. Defaults to False.
 Returns
The penalty loss, detached mean path tensor, and current path length.
 Return type
tuple[Tensor]
 mmagic.models.losses.gradient_penalty_loss(discriminator: torch.nn.Module, real_data: torch.Tensor, fake_data: torch.Tensor, mask: Optional[torch.Tensor] = None, norm_mode: str = 'pixel') torch.Tensor [source]¶
Calculate gradient penalty for wgangp.
 Parameters
discriminator (nn.Module) – Network for the discriminator.
real_data (Tensor) – Real input data.
fake_data (Tensor) – Fake input data.
mask (Tensor) – Masks for inpainting. Default: None.
 Returns
A tensor for gradient penalty.
 Return type
Tensor
 mmagic.models.losses.r1_gradient_penalty_loss(discriminator: torch.nn.Module, real_data: torch.Tensor, mask: Optional[torch.Tensor] = None, norm_mode: str = 'pixel', loss_scaler: Optional[torch.cuda.amp.grad_scaler.GradScaler] = None, use_apex_amp: bool = False) torch.Tensor [source]¶
Calculate R1 gradient penalty for WGANGP.
R1 regularizer comes from: “Which Training Methods for GANs do actually Converge?” ICML’2018
Different from original gradient penalty, this regularizer only penalized gradient w.r.t. real data.
 Parameters
discriminator (nn.Module) – Network for the discriminator.
real_data (Tensor) – Real input data.
mask (Tensor) – Masks for inpainting. Default: None.
norm_mode (str) – This argument decides along which dimension the norm of the gradients will be calculated. Currently, we support [“pixel” , “HWC”]. Defaults to “pixel”.
 Returns
A tensor for gradient penalty.
 Return type
Tensor
 class mmagic.models.losses.GradientLoss(loss_weight: float = 1.0, reduction: str = 'mean')[source]¶
Bases:
torch.nn.Module
Gradient loss.
 Parameters
loss_weight (float) – Loss weight for L1 loss. Default: 1.0.
reduction (str) – Specifies the reduction to apply to the output. Supported choices are ‘none’  ‘mean’  ‘sum’. Default: ‘mean’.
 forward(pred: torch.Tensor, target: torch.Tensor, weight: Optional[torch.Tensor] = None) torch.Tensor [source]¶
 Parameters
pred (Tensor) – of shape (N, C, H, W). Predicted tensor.
target (Tensor) – of shape (N, C, H, W). Ground truth tensor.
weight (Tensor, optional) – of shape (N, C, H, W). Elementwise weights. Default: None.
 class mmagic.models.losses.CLIPLossComps(loss_weight: float = 1.0, data_info: Optional[dict] = None, clip_model: dict = dict(), loss_name: str = 'loss_clip')[source]¶
Bases:
torch.nn.Module
Clip loss. In styleclip, this loss is used to optimize the latent code to generate image that match the text.
In this loss, we may need to provide
image
,text
. Thus, an example of thedata_info
is:1data_info = dict( 2 image='fake_imgs', 3 text='descriptions')
Then, the module will automatically construct this mapping from the input data dictionary.
 Parameters
loss_weight (float, optional) – Weight of this loss item. Defaults to
1.
.data_info (dict, optional) – Dictionary contains the mapping between loss input args and data dictionary. If
None
, this module will directly pass the input data to the loss function. Defaults to None.clip_model (dict, optional) – Kwargs for clip loss model. Defaults to dict().
loss_name (str, optional) – Name of the loss item. If you want this loss item to be included into the backward graph, loss_ must be the prefix of the name. Defaults to ‘loss_clip’.
 forward(*args, **kwargs) torch.Tensor [source]¶
Forward function.
If
self.data_info
is notNone
, a dictionary containing all of the data and necessary modules should be passed into this function. If this dictionary is given as a nonkeyword argument, it should be offered as the first argument. If you are using keyword argument, please name it as outputs_dict.If
self.data_info
isNone
, the input argument or keyword argument will be directly passed to loss function,third_party_net_loss
.
 static loss_name() str [source]¶
Loss Name.
This function must be implemented and will return the name of this loss function. This name will be used to combine different loss items by simple sum operation. In addition, if you want this loss item to be included into the backward graph, loss_ must be the prefix of the name.
 Returns
The name of this loss item.
 Return type
str
 class mmagic.models.losses.DiscShiftLossComps(loss_weight: float = 1.0, data_info: Optional[dict] = None, loss_name: str = 'loss_disc_shift')[source]¶
Bases:
torch.nn.Module
Disc Shift Loss.
This loss is proposed in PGGAN as an auxiliary loss for discriminator.
Note for the design of ``data_info``: In
MMagic
, almost all of loss modules contain the argumentdata_info
, which can be used for constructing the link between the input items (needed in loss calculation) and the data from the generative model. For example, in the training of GAN model, we will collect all of important data/modules into a dictionary:1data_dict_ = dict( 2 gen=self.generator, 3 disc=self.discriminator, 4 disc_pred_fake=disc_pred_fake, 5 disc_pred_real=disc_pred_real, 6 fake_imgs=fake_imgs, 7 real_imgs=real_imgs, 8 iteration=curr_iter, 9 batch_size=batch_size)
But in this loss, we will need to provide
pred
as input. Thus, an example of thedata_info
is:1data_info = dict( 2 pred='disc_pred_fake')
Then, the module will automatically construct this mapping from the input data dictionary.
In addition, in general,
disc_shift_loss
will be applied over real and fake data. In this case, users just need to add this loss module twice, but with differentdata_info
. Our model will automatically add these two items. Parameters
loss_weight (float, optional) – Weight of this loss item. Defaults to
1.
.data_info (dict, optional) – Dictionary contains the mapping between loss input args and data dictionary. If
None
, this module will directly pass the input data to the loss function. Defaults to None.loss_name (str, optional) – Name of the loss item. If you want this loss item to be included into the backward graph, loss_ must be the prefix of the name. Defaults to ‘loss_disc_shift’.
 forward(*args, **kwargs) torch.Tensor [source]¶
Forward function.
If
self.data_info
is notNone
, a dictionary containing all of the data and necessary modules should be passed into this function. If this dictionary is given as a nonkeyword argument, it should be offered as the first argument. If you are using keyword argument, please name it as outputs_dict.If
self.data_info
isNone
, the input argument or keyword argument will be directly passed to loss function,disc_shift_loss
.
 loss_name() str [source]¶
Loss Name.
This function must be implemented and will return the name of this loss function. This name will be used to combine different loss items by simple sum operation. In addition, if you want this loss item to be included into the backward graph, loss_ must be the prefix of the name.
 Returns
The name of this loss item.
 Return type
str
 class mmagic.models.losses.FaceIdLossComps(loss_weight: float = 1.0, data_info: Optional[dict] = None, facenet: dict = dict(type='ArcFace', ir_se50_weights=None), loss_name: str = 'loss_id')[source]¶
Bases:
torch.nn.Module
Face similarity loss. Generally this loss is used to keep the id consistency of the input face image and output face image.
In this loss, we may need to provide
gt
,pred
andx
. Thus, an example of thedata_info
is:1data_info = dict( 2 gt='real_imgs', 3 pred='fake_imgs')
Then, the module will automatically construct this mapping from the input data dictionary.
 Parameters
loss_weight (float, optional) – Weight of this loss item. Defaults to
1.
.data_info (dict, optional) – Dictionary contains the mapping between loss input args and data dictionary. If
None
, this module will directly pass the input data to the loss function. Defaults to None.facenet (dict, optional) – Config dict for facenet. Defaults to dict(type=’ArcFace’, ir_se50_weights=None).
loss_name (str, optional) – Name of the loss item. If you want this loss item to be included into the backward graph, loss_ must be the prefix of the name. Defaults to ‘loss_id’.
 forward(*args, **kwargs) torch.Tensor [source]¶
Forward function.
If
self.data_info
is notNone
, a dictionary containing all of the data and necessary modules should be passed into this function. If this dictionary is given as a nonkeyword argument, it should be offered as the first argument. If you are using keyword argument, please name it as outputs_dict.If
self.data_info
isNone
, the input argument or keyword argument will be directly passed to loss function,third_party_net_loss
.
 loss_name() str [source]¶
Loss Name.
This function must be implemented and will return the name of this loss function. This name will be used to combine different loss items by simple sum operation. In addition, if you want this loss item to be included into the backward graph, loss_ must be the prefix of the name.
 Returns
The name of this loss item.
 Return type
str
 class mmagic.models.losses.GANLossComps(gan_type: str, real_label_val: float = 1.0, fake_label_val: float = 0.0, loss_weight: float = 1.0)[source]¶
Bases:
torch.nn.Module
Define GAN loss.
 Parameters
gan_type (str) – Support ‘vanilla’, ‘lsgan’, ‘wgan’, ‘hinge’, ‘wganlogisticns’.
real_label_val (float) – The value for real label. Default: 1.0.
fake_label_val (float) – The value for fake label. Default: 0.0.
loss_weight (float) – Loss weight. Default: 1.0. Note that loss_weight is only for generators; and it is always 1.0 for discriminators.
 _wgan_loss(input: torch.Tensor, target: bool) torch.Tensor [source]¶
wgan loss.
 Parameters
input (Tensor) – Input tensor.
target (bool) – Target label.
 Returns
wgan loss.
 Return type
Tensor
 _wgan_logistic_ns_loss(input: torch.Tensor, target: bool) torch.Tensor [source]¶
WGAN loss in logistically nonsaturating mode.
This loss is widely used in StyleGANv2.
 Parameters
input (Tensor) – Input tensor.
target (bool) – Target label.
 Returns
wgan loss.
 Return type
Tensor
 get_target_label(input: torch.Tensor, target_is_real: bool) Union[bool, torch.Tensor] [source]¶
Get target label.
 Parameters
input (Tensor) – Input tensor.
target_is_real (bool) – Whether the target is real or fake.
 Returns
Target tensor. Return bool for wgan, otherwise, return Tensor.
 Return type
(bool  Tensor)
 forward(input: torch.Tensor, target_is_real: bool, is_disc: bool = False) torch.Tensor [source]¶
 Parameters
input (Tensor) – The input for the loss module, i.e., the network prediction.
target_is_real (bool) – Whether the targe is real or fake.
is_disc (bool) – Whether the loss for discriminators or not. Default: False.
 Returns
GAN loss value.
 Return type
Tensor
 class mmagic.models.losses.GeneratorPathRegularizerComps(loss_weight: float = 1.0, pl_batch_shrink: int = 1, decay: float = 0.01, pl_batch_size: Optional[int] = None, sync_mean_buffer: bool = False, interval: int = 1, data_info: Optional[dict] = None, use_apex_amp: bool = False, loss_name: str = 'loss_path_regular')[source]¶
Bases:
torch.nn.Module
Generator Path Regularizer.
Path regularization is proposed in StyleGAN2, which can help the improve the continuity of the latent space. More details can be found in: Analyzing and Improving the Image Quality of StyleGAN, CVPR2020.
Users can achieve lazy regularization by setting
interval
arguments here.Note for the design of ``data_info``: In
MMagic
, almost all of loss modules contain the argumentdata_info
, which can be used for constructing the link between the input items (needed in loss calculation) and the data from the generative model. For example, in the training of GAN model, we will collect all of important data/modules into a dictionary:1data_dict_ = dict( 2 gen=self.generator, 3 disc=self.discriminator, 4 fake_imgs=fake_imgs, 5 disc_pred_fake_g=disc_pred_fake_g, 6 iteration=curr_iter, 7 batch_size=batch_size)
But in this loss, we will need to provide
generator
andnum_batches
as input. Thus an example of thedata_info
is:1data_info = dict( 2 generator='gen', 3 num_batches='batch_size')
Then, the module will automatically construct this mapping from the input data dictionary.
 Parameters
loss_weight (float, optional) – Weight of this loss item. Defaults to
1.
.pl_batch_shrink (int, optional) – The factor of shrinking the batch size for saving GPU memory. Defaults to 1.
decay (float, optional) – Decay for moving average of mean path length. Defaults to 0.01.
pl_batch_size (int  None, optional) – The batch size in calculating generator path. Once this argument is set, the
num_batches
will be overridden with this argument and won’t be affected bypl_batch_shrink
. Defaults to None.sync_mean_buffer (bool, optional) – Whether to sync mean path length across all of GPUs. Defaults to False.
interval (int, optional) – The interval of calculating this loss. This argument is used to support lazy regularization. Defaults to 1.
data_info (dict, optional) – Dictionary contains the mapping between loss input args and data dictionary. If
None
, this module will directly pass the input data to the loss function. Defaults to None.loss_name (str, optional) – Name of the loss item. If you want this loss item to be included into the backward graph, loss_ must be the prefix of the name. Defaults to ‘loss_path_regular’.
 forward(*args, **kwargs) torch.Tensor [source]¶
Forward function.
If
self.data_info
is notNone
, a dictionary containing all of the data and necessary modules should be passed into this function. If this dictionary is given as a nonkeyword argument, it should be offered as the first argument. If you are using keyword argument, please name it as outputs_dict.If
self.data_info
isNone
, the input argument or keyword argument will be directly passed to loss function,gen_path_regularizer
.
 loss_name() str [source]¶
Loss Name.
This function must be implemented and will return the name of this loss function. This name will be used to combine different loss items by simple sum operation. In addition, if you want this loss item to be included into the backward graph, loss_ must be the prefix of the name.
 Returns
The name of this loss item.
 Return type
str
 class mmagic.models.losses.GradientPenaltyLossComps(loss_weight: float = 1.0, norm_mode: str = 'pixel', data_info: Optional[dict] = None, loss_name: str = 'loss_gp')[source]¶
Bases:
torch.nn.Module
Gradient Penalty for WGANGP.
In the detailed implementation, there are two streams where one uses the pixelwise gradient norm, but the other adopts normalization along instance (HWC) dimensions. Thus,
norm_mode
are offered to define which mode you want.Note for the design of ``data_info``: In
MMagic
, almost all of loss modules contain the argumentdata_info
, which can be used for constructing the link between the input items (needed in loss calculation) and the data from the generative model. For example, in the training of GAN model, we will collect all of important data/modules into a dictionary:1data_dict_ = dict( 2 gen=self.generator, 3 disc=self.discriminator, 4 disc_pred_fake=disc_pred_fake, 5 disc_pred_real=disc_pred_real, 6 fake_imgs=fake_imgs, 7 real_imgs=real_imgs, 8 iteration=curr_iter, 9 batch_size=batch_size)
But in this loss, we will need to provide
discriminator
,real_data
, andfake_data
as input. Thus, an example of thedata_info
is:1data_info = dict( 2 discriminator='disc', 3 real_data='real_imgs', 4 fake_data='fake_imgs')
Then, the module will automatically construct this mapping from the input data dictionary.
 Parameters
loss_weight (float, optional) – Weight of this loss item. Defaults to
1.
.data_info (dict, optional) – Dictionary contains the mapping between loss input args and data dictionary. If
None
, this module will directly pass the input data to the loss function. Defaults to None.norm_mode (str) – This argument decides along which dimension the norm of the gradients will be calculated. Currently, we support [“pixel” , “HWC”]. Defaults to “pixel”.
loss_name (str, optional) – Name of the loss item. If you want this loss item to be included into the backward graph, loss_ must be the prefix of the name. Defaults to ‘loss_gp’.
 forward(*args, **kwargs) torch.Tensor [source]¶
Forward function.
If
self.data_info
is notNone
, a dictionary containing all of the data and necessary modules should be passed into this function. If this dictionary is given as a nonkeyword argument, it should be offered as the first argument. If you are using keyword argument, please name it as outputs_dict.If
self.data_info
isNone
, the input argument or keyword argument will be directly passed to loss function,gradient_penalty_loss
.
 loss_name() str [source]¶
Loss Name.
This function must be implemented and will return the name of this loss function. This name will be used to combine different loss items by simple sum operation. In addition, if you want this loss item to be included into the backward graph, loss_ must be the prefix of the name.
 Returns
The name of this loss item.
 Return type
str
 class mmagic.models.losses.R1GradientPenaltyComps(loss_weight: float = 1.0, norm_mode: str = 'pixel', interval: int = 1, data_info: Optional[dict] = None, use_apex_amp: bool = False, loss_name: str = 'loss_r1_gp')[source]¶
Bases:
torch.nn.Module
R1 gradient penalty for WGANGP.
R1 regularizer comes from: “Which Training Methods for GANs do actually Converge?” ICML’2018
Different from original gradient penalty, this regularizer only penalized gradient w.r.t. real data.
Note for the design of ``data_info``: In
MMagic
, almost all of loss modules contain the argumentdata_info
, which can be used for constructing the link between the input items (needed in loss calculation) and the data from the generative model. For example, in the training of GAN model, we will collect all of important data/modules into a dictionary:1data_dict_ = dict( 2 gen=self.generator, 3 disc=self.discriminator, 4 disc_pred_fake=disc_pred_fake, 5 disc_pred_real=disc_pred_real, 6 fake_imgs=fake_imgs, 7 real_imgs=real_imgs, 8 iteration=curr_iter, 9 batch_size=batch_size)
But in this loss, we will need to provide
discriminator
andreal_data
as input. Thus, an example of thedata_info
is:1data_info = dict( 2 discriminator='disc', 3 real_data='real_imgs')
Then, the module will automatically construct this mapping from the input data dictionary.
 Parameters
loss_weight (float, optional) – Weight of this loss item. Defaults to
1.
.data_info (dict, optional) – Dictionary contains the mapping between loss input args and data dictionary. If
None
, this module will directly pass the input data to the loss function. Defaults to None.norm_mode (str) – This argument decides along which dimension the norm of the gradients will be calculated. Currently, we support [“pixel” , “HWC”]. Defaults to “pixel”.
interval (int, optional) – The interval of calculating this loss. Defaults to 1.
loss_name (str, optional) – Name of the loss item. If you want this loss item to be included into the backward graph, loss_ must be the prefix of the name. Defaults to ‘loss_r1_gp’.
 forward(*args, **kwargs) torch.Tensor [source]¶
Forward function.
If
self.data_info
is notNone
, a dictionary containing all of the data and necessary modules should be passed into this function. If this dictionary is given as a nonkeyword argument, it should be offered as the first argument. If you are using keyword argument, please name it as outputs_dict.If
self.data_info
isNone
, the input argument or keyword argument will be directly passed to loss function,r1_gradient_penalty_loss
.
 loss_name() str [source]¶
Loss Name.
This function must be implemented and will return the name of this loss function. This name will be used to combine different loss items by simple sum operation. In addition, if you want this loss item to be included into the backward graph, loss_ must be the prefix of the name.
 Returns
The name of this loss item.
 Return type
str
 mmagic.models.losses.mask_reduce_loss(loss: torch.Tensor, weight: Optional[torch.Tensor] = None, reduction: str = 'mean', sample_wise: bool = False) torch.Tensor [source]¶
Apply elementwise weight and reduce loss.
 Parameters
loss (Tensor) – Elementwise loss.
weight (Tensor) – Elementwise weights. Default: None.
reduction (str) – Same as builtin losses of PyTorch. Options are “none”, “mean” and “sum”. Default: ‘mean’.
sample_wise (bool) – Whether calculate the loss samplewise. This argument only takes effect when reduction is ‘mean’ and weight (argument of forward()) is not None. It will first reduces loss with ‘mean’ persample, and then it means over all the samples. Default: False.
 Returns
Processed loss values.
 Return type
Tensor
 mmagic.models.losses.reduce_loss(loss: torch.Tensor, reduction: str) torch.Tensor [source]¶
Reduce loss as specified.
 Parameters
loss (Tensor) – Elementwise loss tensor.
reduction (str) – Options are “none”, “mean” and “sum”.
 Returns
Reduced loss tensor.
 Return type
Tensor
 class mmagic.models.losses.PerceptualLoss(layer_weights: dict, layer_weights_style: Optional[dict] = None, vgg_type: str = 'vgg19', use_input_norm: bool = True, perceptual_weight: float = 1.0, style_weight: float = 1.0, norm_img: bool = True, pretrained: str = 'torchvision://vgg19', criterion: str = 'l1')[source]¶
Bases:
torch.nn.Module
Perceptual loss with commonly used style loss.
 Parameters
layers_weights (dict) – The weight for each layer of vgg feature for perceptual loss. Here is an example: {‘4’: 1., ‘9’: 1., ‘18’: 1.}, which means the 5th, 10th and 18th feature layer will be extracted with weight 1.0 in calculating losses.
layers_weights_style (dict) – The weight for each layer of vgg feature for style loss. If set to ‘None’, the weights are set equal to the weights for perceptual loss. Default: None.
vgg_type (str) – The type of vgg network used as feature extractor. Default: ‘vgg19’.
use_input_norm (bool) – If True, normalize the input image in vgg. Default: True.
perceptual_weight (float) – If perceptual_weight > 0, the perceptual loss will be calculated and the loss will multiplied by the weight. Default: 1.0.
style_weight (float) – If style_weight > 0, the style loss will be calculated and the loss will multiplied by the weight. Default: 1.0.
norm_img (bool) – If True, the image will be normed to [0, 1]. Note that this is different from the use_input_norm which norm the input in in forward function of vgg according to the statistics of dataset. Importantly, the input image must be in range [1, 1].
pretrained (str) – Path for pretrained weights. Default: ‘torchvision://vgg19’.
criterion (str) – Criterion type. Options are ‘l1’ and ‘mse’. Default: ‘l1’.
 class mmagic.models.losses.PerceptualVGG(layer_name_list: List[str], vgg_type: str = 'vgg19', use_input_norm: bool = True, pretrained: str = 'torchvision://vgg19')[source]¶
Bases:
torch.nn.Module
VGG network used in calculating perceptual loss.
In this implementation, we allow users to choose whether use normalization in the input feature and the type of vgg network. Note that the pretrained path must fit the vgg type.
 Parameters
layer_name_list (list[str]) – According to the name in this list, forward function will return the corresponding features. This list contains the name each layer in vgg.feature. An example of this list is [‘4’, ‘10’].
vgg_type (str) – Set the type of vgg network. Default: ‘vgg19’.
use_input_norm (bool) – If True, normalize the input image. Importantly, the input feature must in the range [0, 1]. Default: True.
pretrained (str) – Path for pretrained weights. Default: ‘torchvision://vgg19’
 class mmagic.models.losses.TransferalPerceptualLoss(loss_weight: float = 1.0, use_attention: bool = True, criterion: str = 'mse')[source]¶
Bases:
torch.nn.Module
Transferal perceptual loss.
 Parameters
loss_weight (float) – Loss weight. Default: 1.0.
use_attention (bool) – If True, use softattention tensor. Default: True
criterion (str) – Criterion type. Options are ‘l1’ and ‘mse’. Default: ‘mse’.
 forward(maps: Tuple[torch.Tensor], soft_attention: torch.Tensor, textures: Tuple[torch.Tensor]) torch.Tensor [source]¶
Forward function.
 Parameters
maps (Tuple[Tensor]) – Input tensors.
soft_attention (Tensor) – Softattention tensor.
textures (Tuple[Tensor]) – Groundtruth tensors.
 Returns
Forward results.
 Return type
Tensor
 class mmagic.models.losses.CharbonnierLoss(loss_weight: float = 1.0, reduction: str = 'mean', sample_wise: bool = False, eps: float = 1e12)[source]¶
Bases:
torch.nn.Module
Charbonnier loss (one variant of Robust L1Loss, a differentiable variant of L1Loss).
 Described in “Deep Laplacian Pyramid Networks for Fast and Accurate
SuperResolution”.
 Parameters
loss_weight (float) – Loss weight for L1 loss. Default: 1.0.
reduction (str) – Specifies the reduction to apply to the output. Supported choices are ‘none’  ‘mean’  ‘sum’. Default: ‘mean’.
sample_wise (bool) – Whether calculate the loss samplewise. This argument only takes effect when reduction is ‘mean’ and weight (argument of forward()) is not None. It will first reduces loss with ‘mean’ persample, and then it means over all the samples. Default: False.
eps (float) – A value used to control the curvature near zero. Default: 1e12.
 forward(pred: torch.Tensor, target: torch.Tensor, weight: Optional[torch.Tensor] = None, **kwargs) torch.Tensor [source]¶
Forward Function.
 Parameters
pred (Tensor) – of shape (N, C, H, W). Predicted tensor.
target (Tensor) – of shape (N, C, H, W). Ground truth tensor.
weight (Tensor, optional) – of shape (N, C, H, W). Elementwise weights. Default: None.
 class mmagic.models.losses.L1Loss(loss_weight: float = 1.0, reduction: str = 'mean', sample_wise: bool = False)[source]¶
Bases:
torch.nn.Module
L1 (mean absolute error, MAE) loss.
 Parameters
loss_weight (float) – Loss weight for L1 loss. Default: 1.0.
reduction (str) – Specifies the reduction to apply to the output. Supported choices are ‘none’  ‘mean’  ‘sum’. Default: ‘mean’.
sample_wise (bool) – Whether calculate the loss samplewise. This argument only takes effect when reduction is ‘mean’ and weight (argument of forward()) is not None. It will first reduce loss with ‘mean’ persample, and then it means over all the samples. Default: False.
 forward(pred: torch.Tensor, target: torch.Tensor, weight: Optional[torch.Tensor] = None, **kwargs) torch.Tensor [source]¶
Forward Function.
 Parameters
pred (Tensor) – of shape (N, C, H, W). Predicted tensor.
target (Tensor) – of shape (N, C, H, W). Ground truth tensor.
weight (Tensor, optional) – of shape (N, C, H, W). Elementwise weights. Default: None.
 class mmagic.models.losses.MaskedTVLoss(loss_weight: float = 1.0)[source]¶
Bases:
L1Loss
Masked TV loss.
 Parameters
loss_weight (float, optional) – Loss weight. Defaults to 1.0.
 class mmagic.models.losses.MSELoss(loss_weight: float = 1.0, reduction: str = 'mean', sample_wise: bool = False)[source]¶
Bases:
torch.nn.Module
MSE (L2) loss.
 Parameters
loss_weight (float) – Loss weight for MSE loss. Default: 1.0.
reduction (str) – Specifies the reduction to apply to the output. Supported choices are ‘none’  ‘mean’  ‘sum’. Default: ‘mean’.
sample_wise (bool) – Whether calculate the loss samplewise. This argument only takes effect when reduction is ‘mean’ and weight (argument of forward()) is not None. It will first reduces loss with ‘mean’ persample, and then it means over all the samples. Default: False.
 forward(pred: torch.Tensor, target: torch.Tensor, weight: Optional[torch.Tensor] = None, **kwargs) torch.Tensor [source]¶
Forward Function.
 Parameters
pred (Tensor) – of shape (N, C, H, W). Predicted tensor.
target (Tensor) – of shape (N, C, H, W). Ground truth tensor.
weight (Tensor, optional) – of shape (N, C, H, W). Elementwise weights. Default: None.
 class mmagic.models.losses.PSNRLoss(loss_weight: float = 1.0, toY: bool = False)[source]¶
Bases:
torch.nn.Module
PSNR Loss in “HINet: Half Instance Normalization Network for Image Restoration”.
 Parameters
loss_weight (float, optional) – Loss weight. Defaults to 1.0.
reduction – reduction for PSNR. Can only be mean here.
toY – change to calculate the PSNR of Y channel in YCbCr format