Shortcuts

mmedit.models.losses

Package Contents

Classes

CLIPLoss

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

CharbonnierCompLoss

Charbonnier composition loss.

L1CompositionLoss

L1 composition loss.

MSECompositionLoss

MSE (L2) composition loss.

FaceIdLoss

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

LightCNNFeatureLoss

Feature loss of DICGAN, based on LightCNN.

DiscShiftLoss

Disc shift loss.

GANLoss

Define GAN loss.

GaussianBlur

A Gaussian filter which blurs a given tensor with a two-dimensional

GradientPenaltyLoss

Gradient penalty loss for wgan-gp.

GradientLoss

Gradient loss.

CLIPLossComps

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

DiscShiftLossComps

Disc Shift Loss.

FaceIdLossComps

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

GANLossComps

Define GAN loss.

GeneratorPathRegularizerComps

Generator Path Regularizer.

GradientPenaltyLossComps

Gradient Penalty for WGAN-GP.

R1GradientPenaltyComps

R1 gradient penalty for WGAN-GP.

PerceptualLoss

Perceptual loss with commonly used style loss.

PerceptualVGG

VGG network used in calculating perceptual loss.

TransferalPerceptualLoss

Transferal perceptual loss.

CharbonnierLoss

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

L1Loss

L1 (mean absolute error, MAE) loss.

MaskedTVLoss

Masked TV loss.

MSELoss

MSE (L2) loss.

PSNRLoss

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

Functions

disc_shift_loss(pred)

Disc Shift loss.

gen_path_regularizer(generator, num_batches, ...[, ...])

Generator Path Regularization.

gradient_penalty_loss(discriminator, real_data, fake_data)

Calculate gradient penalty for wgan-gp.

r1_gradient_penalty_loss(discriminator, real_data[, ...])

Calculate R1 gradient penalty for WGAN-GP.

mask_reduce_loss(loss[, weight, reduction, sample_wise])

Apply element-wise weight and reduce loss.

reduce_loss(loss, reduction)

Reduce loss as specified.

tv_loss(input)

L2 total variation loss, as in Mahendran et al.

class mmedit.models.losses.CLIPLoss(loss_weight=1.0, data_info=None, clip_model=dict(), loss_name='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 the data_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, text)

Forward function.

If self.data_info is not None, a dictionary containing all of the data and necessary modules should be passed into this function. If this dictionary is given as a non-keyword 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 is None, the input argument or key-word argument will be directly passed to loss function, third_party_net_loss.

class mmedit.models.losses.CharbonnierCompLoss(loss_weight=1.0, reduction='mean', sample_wise=False, eps=1e-12)[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 sample-wise. This argument only takes effect when reduction is ‘mean’ and weight (argument of forward()) is not None. It will first reduces loss with ‘mean’ per-sample, and then it means over all the samples. Default: False.

  • eps (float) – A value used to control the curvature near zero. Default: 1e-12.

forward(pred_alpha, fg, bg, ori_merged, weight=None, **kwargs)
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 mmedit.models.losses.L1CompositionLoss(loss_weight=1.0, reduction='mean', sample_wise=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 sample-wise. This argument only takes effect when reduction is ‘mean’ and weight (argument of forward()) is not None. It will first reduces loss with ‘mean’ per-sample, and then it means over all the samples. Default: False.

forward(pred_alpha, fg, bg, ori_merged, weight=None, **kwargs)
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 mmedit.models.losses.MSECompositionLoss(loss_weight=1.0, reduction='mean', sample_wise=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 sample-wise. This argument only takes effect when reduction is ‘mean’ and weight (argument of forward()) is not None. It will first reduces loss with ‘mean’ per-sample, and then it means over all the samples. Default: False.

forward(pred_alpha, fg, bg, ori_merged, weight=None, **kwargs)
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 mmedit.models.losses.FaceIdLoss(loss_weight=1.0, data_info=None, facenet=dict(type='ArcFace', ir_se50_weights=None), loss_name='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 and x. Thus, an example of the data_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’.

forward(pred=None, gt=None)

Forward function.

class mmedit.models.losses.LightCNNFeatureLoss(pretrained, loss_weight=1.0, criterion='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’.

forward(pred, gt)

Forward function.

Parameters
  • pred (Tensor) – Predicted tensor.

  • gt (Tensor) – GT tensor.

Returns

Forward results.

Return type

Tensor

class mmedit.models.losses.DiscShiftLoss(loss_weight=0.1)[source]

Bases: torch.nn.Module

Disc shift loss.

Parameters

loss_weight (float, optional) – Loss weight. Defaults to 1.0.

forward(x)

Forward function.

Parameters

x (Tensor) – Tensor with shape (n, c, h, w)

Returns

Loss.

Return type

Tensor

class mmedit.models.losses.GANLoss(gan_type, real_label_val=1.0, fake_label_val=0.0, loss_weight=1.0)[source]

Bases: torch.nn.Module

Define GAN loss.

Parameters
  • gan_type (str) – Support ‘vanilla’, ‘lsgan’, ‘wgan’, ‘hinge’.

  • 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, target)

wgan loss.

Parameters
  • input (Tensor) – Input tensor.

  • target (bool) – Target label.

Returns

wgan loss.

Return type

Tensor

get_target_label(input, target_is_real)

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, target_is_real, is_disc=False, mask=None)
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 mmedit.models.losses.GaussianBlur(kernel_size=(71, 71), sigma=(10.0, 10.0))[source]

Bases: torch.nn.Module

A Gaussian filter which blurs a given tensor with a two-dimensional 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 Gaussian-blurred 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)

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, sigma)

Get the two-dimensional 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, sigma)

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, sigma)

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

forward(x)

Forward function.

Parameters

x (Tensor) – Tensor with shape (n, c, h, w)

Returns

The Gaussian-blurred tensor.

Return type

Tensor

class mmedit.models.losses.GradientPenaltyLoss(loss_weight=1.0)[source]

Bases: torch.nn.Module

Gradient penalty loss for wgan-gp.

Parameters

loss_weight (float) – Loss weight. Default: 1.0.

forward(discriminator, real_data, fake_data, mask=None)

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

mmedit.models.losses.disc_shift_loss(pred)[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

mmedit.models.losses.gen_path_regularizer(generator, num_batches, mean_path_length, pl_batch_shrink=1, decay=0.01, weight=1.0, pl_batch_size=None, sync_mean_buffer=False, loss_scaler=None, use_apex_amp=False)[source]

Generator Path Regularization.

Path regularization is proposed in StyelGAN2, 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 affectted by pl_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]

mmedit.models.losses.gradient_penalty_loss(discriminator, real_data, fake_data, mask=None, norm_mode='pixel')[source]

Calculate gradient penalty for wgan-gp.

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

mmedit.models.losses.r1_gradient_penalty_loss(discriminator, real_data, mask=None, norm_mode='pixel', loss_scaler=None, use_apex_amp=False)[source]

Calculate R1 gradient penalty for WGAN-GP.

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 mmedit.models.losses.GradientLoss(loss_weight=1.0, reduction='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, target, weight=None)
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). Element-wise weights. Default: None.

class mmedit.models.losses.CLIPLossComps(loss_weight=1.0, data_info=None, clip_model=dict(), loss_name='loss_clip')

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 the data_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)

Forward function.

If self.data_info is not None, a dictionary containing all of the data and necessary modules should be passed into this function. If this dictionary is given as a non-keyword 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 is None, the input argument or key-word argument will be directly passed to loss function, third_party_net_loss.

static loss_name()

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 mmedit.models.losses.DiscShiftLossComps(loss_weight=1.0, data_info=None, loss_name='loss_disc_shift')

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 MMEditing, almost all of loss modules contain the argument data_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:

Code from StaticUnconditionalGAN, train_step
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 the data_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 different data_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)

Forward function.

If self.data_info is not None, a dictionary containing all of the data and necessary modules should be passed into this function. If this dictionary is given as a non-keyword 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 is None, the input argument or key-word argument will be directly passed to loss function, disc_shift_loss.

loss_name()

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 mmedit.models.losses.FaceIdLossComps(loss_weight=1.0, data_info=None, facenet=dict(type='ArcFace', ir_se50_weights=None), loss_name='loss_id')

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 and x. Thus, an example of the data_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)

Forward function.

If self.data_info is not None, a dictionary containing all of the data and necessary modules should be passed into this function. If this dictionary is given as a non-keyword 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 is None, the input argument or key-word argument will be directly passed to loss function, third_party_net_loss.

loss_name()

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 mmedit.models.losses.GANLossComps(gan_type, real_label_val=1.0, fake_label_val=0.0, loss_weight=1.0)

Bases: torch.nn.Module

Define GAN loss.

Parameters
  • gan_type (str) – Support ‘vanilla’, ‘lsgan’, ‘wgan’, ‘hinge’, ‘wgan-logistic-ns’.

  • 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, target)

wgan loss.

Parameters
  • input (Tensor) – Input tensor.

  • target (bool) – Target label.

Returns

wgan loss.

Return type

Tensor

_wgan_logistic_ns_loss(input, target)

WGAN loss in logistically non-saturating 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, target_is_real)

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, target_is_real, is_disc=False)
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 mmedit.models.losses.GeneratorPathRegularizerComps(loss_weight=1.0, pl_batch_shrink=1, decay=0.01, pl_batch_size=None, sync_mean_buffer=False, interval=1, data_info=None, use_apex_amp=False, loss_name='loss_path_regular')

Bases: torch.nn.Module

Generator Path Regularizer.

Path regularization is proposed in StyelGAN2, 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 MMEditing, almost all of loss modules contain the argument data_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:

Code from StaticUnconditionalGAN, train_step
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 and num_batches as input. Thus an example of the data_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 affectted by pl_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)

Forward function.

If self.data_info is not None, a dictionary containing all of the data and necessary modules should be passed into this function. If this dictionary is given as a non-keyword 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 is None, the input argument or key-word argument will be directly passed to loss function, gen_path_regularizer.

loss_name()

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 mmedit.models.losses.GradientPenaltyLossComps(loss_weight=1.0, norm_mode='pixel', data_info=None, loss_name='loss_gp')

Bases: torch.nn.Module

Gradient Penalty for WGAN-GP.

In the detailed implementation, there are two streams where one uses the pixel-wise 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 MMEditing, almost all of loss modules contain the argument data_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:

Code from StaticUnconditionalGAN, train_step
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, and fake_data as input. Thus, an example of the data_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)

Forward function.

If self.data_info is not None, a dictionary containing all of the data and necessary modules should be passed into this function. If this dictionary is given as a non-keyword 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 is None, the input argument or key-word argument will be directly passed to loss function, gradient_penalty_loss.

loss_name()

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 mmedit.models.losses.R1GradientPenaltyComps(loss_weight=1.0, norm_mode='pixel', interval=1, data_info=None, use_apex_amp=False, loss_name='loss_r1_gp')

Bases: torch.nn.Module

R1 gradient penalty for WGAN-GP.

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 MMEditing, almost all of loss modules contain the argument data_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:

Code from StaticUnconditionalGAN, train_step
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 and real_data as input. Thus, an example of the data_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)

Forward function.

If self.data_info is not None, a dictionary containing all of the data and necessary modules should be passed into this function. If this dictionary is given as a non-keyword 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 is None, the input argument or key-word argument will be directly passed to loss function, r1_gradient_penalty_loss.

loss_name()

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

mmedit.models.losses.mask_reduce_loss(loss, weight=None, reduction='mean', sample_wise=False)[source]

Apply element-wise weight and reduce loss.

Parameters
  • loss (Tensor) – Element-wise loss.

  • weight (Tensor) – Element-wise weights. Default: None.

  • reduction (str) – Same as built-in losses of PyTorch. Options are “none”, “mean” and “sum”. Default: ‘mean’.

  • sample_wise (bool) – Whether calculate the loss sample-wise. This argument only takes effect when reduction is ‘mean’ and weight (argument of forward()) is not None. It will first reduces loss with ‘mean’ per-sample, and then it means over all the samples. Default: False.

Returns

Processed loss values.

Return type

Tensor

mmedit.models.losses.reduce_loss(loss, reduction)[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 mmedit.models.losses.PerceptualLoss(layer_weights, layer_weights_style=None, vgg_type='vgg19', use_input_norm=True, perceptual_weight=1.0, style_weight=1.0, norm_img=True, pretrained='torchvision://vgg19', criterion='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’.

forward(x, gt)

Forward function.

Parameters
  • x (Tensor) – Input tensor with shape (n, c, h, w).

  • gt (Tensor) – Ground-truth tensor with shape (n, c, h, w).

Returns

Forward results.

Return type

Tensor

_gram_mat(x)

Calculate Gram matrix.

Parameters

x (torch.Tensor) – Tensor with shape of (n, c, h, w).

Returns

Gram matrix.

Return type

torch.Tensor

class mmedit.models.losses.PerceptualVGG(layer_name_list, vgg_type='vgg19', use_input_norm=True, pretrained='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’

forward(x)

Forward function.

Parameters

x (Tensor) – Input tensor with shape (n, c, h, w).

Returns

Forward results.

Return type

Tensor

init_weights(model, pretrained)

Init weights.

Parameters
  • model (nn.Module) – Models to be inited.

  • pretrained (str) – Path for pretrained weights.

class mmedit.models.losses.TransferalPerceptualLoss(loss_weight=1.0, use_attention=True, criterion='mse')[source]

Bases: torch.nn.Module

Transferal perceptual loss.

Parameters
  • loss_weight (float) – Loss weight. Default: 1.0.

  • use_attention (bool) – If True, use soft-attention tensor. Default: True

  • criterion (str) – Criterion type. Options are ‘l1’ and ‘mse’. Default: ‘mse’.

forward(maps, soft_attention, textures)

Forward function.

Parameters
  • maps (Tuple[Tensor]) – Input tensors.

  • soft_attention (Tensor) – Soft-attention tensor.

  • textures (Tuple[Tensor]) – Ground-truth tensors.

Returns

Forward results.

Return type

Tensor

class mmedit.models.losses.CharbonnierLoss(loss_weight=1.0, reduction='mean', sample_wise=False, eps=1e-12)[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

Super-Resolution”.

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 sample-wise. This argument only takes effect when reduction is ‘mean’ and weight (argument of forward()) is not None. It will first reduces loss with ‘mean’ per-sample, and then it means over all the samples. Default: False.

  • eps (float) – A value used to control the curvature near zero. Default: 1e-12.

forward(pred, target, weight=None, **kwargs)

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). Element-wise weights. Default: None.

class mmedit.models.losses.L1Loss(loss_weight=1.0, reduction='mean', sample_wise=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 sample-wise. This argument only takes effect when reduction is ‘mean’ and weight (argument of forward()) is not None. It will first reduce loss with ‘mean’ per-sample, and then it means over all the samples. Default: False.

forward(pred, target, weight=None, **kwargs)

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). Element-wise weights. Default: None.

class mmedit.models.losses.MaskedTVLoss(loss_weight=1.0)[source]

Bases: L1Loss

Masked TV loss.

Parameters

loss_weight (float, optional) – Loss weight. Defaults to 1.0.

forward(pred, mask=None)

Forward function.

Parameters
  • pred (torch.Tensor) – Tensor with shape of (n, c, h, w).

  • mask (torch.Tensor, optional) – Tensor with shape of (n, 1, h, w). Defaults to None.

Returns

[description]

Return type

[type]

class mmedit.models.losses.MSELoss(loss_weight=1.0, reduction='mean', sample_wise=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 sample-wise. This argument only takes effect when reduction is ‘mean’ and weight (argument of forward()) is not None. It will first reduces loss with ‘mean’ per-sample, and then it means over all the samples. Default: False.

forward(pred, target, weight=None, **kwargs)

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). Element-wise weights. Default: None.

class mmedit.models.losses.PSNRLoss(loss_weight=1.0, toY=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

forward(pred, target)
mmedit.models.losses.tv_loss(input)[source]

L2 total variation loss, as in Mahendran et al.

Read the Docs v: zyh/api-rendering
Versions
master
latest
stable
zyh-doc-notfound-extend
zyh-api-rendering
Downloads
pdf
epub
On Read the Docs
Project Home
Builds

Free document hosting provided by Read the Docs.