imaginaire.losses package

Submodules

imaginaire.losses.dict module

class imaginaire.losses.dict.DictLoss(criterion='l1')

Bases: torch.nn.modules.module.Module

forward(fake, real)

Return the target vector for the l1/l2 loss computation.

Parameters

• fake (dict, list or tuple) – Discriminator features of fake images.

• real (dict, list or tuple) – Discriminator features of real images.

Returns

Loss value.

Return type

loss (tensor)

training = None

imaginaire.losses.feature_matching module

class imaginaire.losses.feature_matching.FeatureMatchingLoss(criterion='l1')

Bases: torch.nn.modules.module.Module

Compute feature matching loss

forward(fake_features, real_features)

Return the target vector for the binary cross entropy loss computation.

Parameters

• fake_features (list of lists) – Discriminator features of fake images.

• real_features (list of lists) – Discriminator features of real images.

Returns

Loss value.

Return type

(tensor)

training = None

imaginaire.losses.flow module

class imaginaire.losses.flow.FlowLoss(cfg)

Bases: torch.nn.modules.module.Module

Flow loss constructor.

Parameters

cfg (obj) – Configuration.

compute_flow_loss(flow, warped_image, tgt_image, flow_gt, flow_conf_gt, fg_mask)

Compute losses on the generated flow map.

Parameters

• flow (tensor) – Generated flow map.

• warped_image (tensor) – Warped image using the flow map.

• tgt_image (tensor) – Target image for the warped image.

• flow_gt (tensor) – Ground truth flow map.

• flow_conf_gt (tensor) – Confidence for the ground truth flow map.

• fg_mask (tensor) – Foreground mask for the target image.

Returns

• loss_flow_L1 (tensor): L1 loss compared to ground truth flow.

• loss_flow_warp (tensor): L1 loss between the warped image and

the target image when using the flow to warp.

Return type

(dict)

compute_flow_losses(flow, warped_images, tgt_image, flow_gt, flow_conf_gt, fg_mask, tgt_label, ref_label)

Compute losses on the generated flow maps.

Parameters

• flow (tensor or list of tensors) – Generated flow maps. warped_images (tensor or list of tensors): Warped images using the flow maps.

• tgt_image (tensor) – Target image for the warped image. flow_gt (tensor or list of tensors): Ground truth flow maps.

• flow_conf_gt (tensor or list of tensors) – Confidence for the ground truth flow maps.

• fg_mask (tensor) – Foreground mask for the target image.

• tgt_label (tensor) – Target label map.

• ref_label (tensor) – Reference label map.

Returns

• loss_flow_L1 (tensor): L1 loss compared to ground truth flow.

• loss_flow_warp (tensor): L1 loss between the warped image and the target image when using the flow to warp.

• body_mask_diff (tensor): Difference between warped body part map and target body part map. Used for pose dataset only.

Return type

(dict)

compute_mask_loss(occ_mask, warped_image, tgt_image)

Compute losses on the generated occlusion mask.

Parameters

• occ_mask (tensor) – Generated occlusion mask.

• warped_image (tensor) – Warped image using the flow map.

• tgt_image (tensor) – Target image for the warped image.

Returns

Loss for the mask.

Return type

(tensor)

compute_mask_losses(occ_mask, fake_image, warped_image, tgt_label, tgt_image, fg_mask, ref_fg_mask, body_mask_diff)

Compute losses on the generated occlusion masks.

Parameters

• occ_mask (tensor or list of tensors) – Generated occlusion masks.

• fake_image (tensor) – Generated image.

• warped_image (tensor or list of tensors) – Warped images using the flow maps.

• tgt_label (tensor) – Target label map.

• tgt_image (tensor) – Target image for the warped image.

• fg_mask (tensor) – Foreground mask for the target image.

• ref_fg_mask (tensor) – Foreground mask for the reference image.

• body_mask_diff (tensor) – Difference between warped body part map

• target body part map. Used for pose dataset only. (and) –

Returns

Loss for the mask.

Return type

(tensor)

forward(data, net_G_output, current_epoch)

Compute losses on the output flow and occlusion mask.

Parameters

• data (dict) – Input data.

• net_G_output (dict) – Generator output.

• current_epoch (int) – Current training epoch number.

Returns

• loss_flow_L1 (tensor): L1 loss compared to ground truth flow.

• loss_flow_warp (tensor): L1 loss between the warped image and the target image when using the flow to warp.

• loss_mask (tensor): Loss for the occlusion mask.

Return type

(dict)

training = None

class imaginaire.losses.flow.MaskedL1Loss(normalize_over_valid=False)

Bases: torch.nn.modules.module.Module

Masked L1 loss constructor.

forward(input, target, mask)

Masked L1 loss computation.

Parameters

• input (tensor) – Input tensor.

• target (tensor) – Target tensor.

• mask (tensor) – Mask to be applied to the output loss.

Returns

Loss value.

Return type

(tensor)

training = None

imaginaire.losses.gan module

class imaginaire.losses.gan.GANLoss(gan_mode, target_real_label=1.0, target_fake_label=0.0, decay_k=1.0, min_k=1.0, separate_topk=False)

Bases: torch.nn.modules.module.Module

GAN loss constructor.

Parameters

• 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.

forward(dis_output, t_real, dis_update=True, reduce=True)

GAN loss computation.

Parameters

• 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 value.

Return type

loss (tensor)

get_target_tensor(dis_output, t_real)

Return the target vector for the binary cross entropy loss computation.

Parameters

• 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 vector.

Return type

target (tensor)

loss(dis_output, t_real, dis_update=True)

GAN loss computation.

Parameters

• 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 value.

Return type

loss (tensor)

topk_anneal()

Anneal k after each epoch.

training = None

imaginaire.losses.info_nce module

class imaginaire.losses.info_nce.GatherLayer

Bases: torch.autograd.function.Function

static backward(ctx, *grads)

Defines a formula for differentiating the operation.

This function is to be overridden by all subclasses.

It must accept a context ctx as the first argument, followed by as many outputs did forward() return, and it should return as many tensors, as there were inputs to forward(). Each argument is the gradient w.r.t the given output, and each returned value should be the gradient w.r.t. the corresponding input.

The context can be used to retrieve tensors saved during the forward pass. It also has an attribute ctx.needs_input_grad as a tuple of booleans representing whether each input needs gradient. E.g., backward() will have ctx.needs_input_grad[0] = True if the first input to forward() needs gradient computated w.r.t. the output.

static forward(ctx, input)

Performs the operation.

This function is to be overridden by all subclasses.

It must accept a context ctx as the first argument, followed by any number of arguments (tensors or other types).

The context can be used to store tensors that can be then retrieved during the backward pass.

class imaginaire.losses.info_nce.InfoNCELoss(temperature=0.07, gather_distributed=True, learn_temperature=True, single_direction=False, flatten=True)

Bases: torch.nn.modules.module.Module

forward(features_a, features_b, gather_distributed=None, eps=1e-08)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

training = None

imaginaire.losses.kl module

class imaginaire.losses.kl.GaussianKLLoss

Bases: torch.nn.modules.module.Module

Compute KL loss in VAE for Gaussian distributions

forward(mu, logvar=None)

Compute loss

Parameters

• mu (tensor) – mean

• logvar (tensor) – logarithm of variance

training = None

imaginaire.losses.perceptual module

class imaginaire.losses.perceptual.PerceptualLoss(network='vgg19', layers='relu_4_1', weights=None, criterion='l1', resize=False, resize_mode='bilinear', num_scales=1, per_sample_weight=False, info_nce_temperature=0.07, info_nce_gather_distributed=True, info_nce_learn_temperature=True, info_nce_flatten=True)

Bases: torch.nn.modules.module.Module

Perceptual loss initialization.

Parameters

• network (str) – The name of the loss network: ‘vgg16’ | ‘vgg19’.

• layers (str or list of str) – The layers used to compute the loss.

• (float or list of float (weights) – The loss weights of each layer.

• criterion (str) – The type of distance function: ‘l1’ | ‘l2’.

• resize (bool) – If True, resize the input images to 224x224.

• resize_mode (str) – Algorithm used for resizing.

• num_scales (int) – The loss will be evaluated at original size and this many times downsampled sizes.

• per_sample_weight (bool) – Output loss for individual samples in the batch instead of mean loss.

forward(inp, target, per_sample_weights=None)

Perceptual loss forward.

Parameters

• inp (4D tensor) – Input tensor.

• target (4D tensor) – Ground truth tensor, same shape as the input.

• per_sample_weight (bool) – Output loss for individual samples in the batch instead of mean loss.

Returns

The perceptual loss.

Return type

(scalar tensor)

training = None

imaginaire.losses.weighted_mse module

class imaginaire.losses.weighted_mse.WeightedMSELoss(reduction='mean')

Bases: torch.nn.modules.module.Module

Compute Weighted MSE loss

forward(input, target, weight)

Return weighted MSE Loss. :param input: :type input: tensor :param target: :type target: tensor :param weight: :type weight: tensor

Returns

Loss value.

Return type

(tensor)

training = None

Module contents

class imaginaire.losses.GANLoss(gan_mode, target_real_label=1.0, target_fake_label=0.0, decay_k=1.0, min_k=1.0, separate_topk=False)

Bases: torch.nn.modules.module.Module

GAN loss constructor.

Parameters

• 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.

forward(dis_output, t_real, dis_update=True, reduce=True)

GAN loss computation.

Parameters

• 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 value.

Return type

loss (tensor)

get_target_tensor(dis_output, t_real)

Return the target vector for the binary cross entropy loss computation.

Parameters

• 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 vector.

Return type

target (tensor)

loss(dis_output, t_real, dis_update=True)

GAN loss computation.

Parameters

• 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 value.

Return type

loss (tensor)

topk_anneal()

Anneal k after each epoch.

training = None

class imaginaire.losses.PerceptualLoss(network='vgg19', layers='relu_4_1', weights=None, criterion='l1', resize=False, resize_mode='bilinear', num_scales=1, per_sample_weight=False, info_nce_temperature=0.07, info_nce_gather_distributed=True, info_nce_learn_temperature=True, info_nce_flatten=True)

Bases: torch.nn.modules.module.Module

Perceptual loss initialization.

Parameters

• network (str) – The name of the loss network: ‘vgg16’ | ‘vgg19’.

• layers (str or list of str) – The layers used to compute the loss.

• (float or list of float (weights) – The loss weights of each layer.

• criterion (str) – The type of distance function: ‘l1’ | ‘l2’.

• resize (bool) – If True, resize the input images to 224x224.

• resize_mode (str) – Algorithm used for resizing.

• num_scales (int) – The loss will be evaluated at original size and this many times downsampled sizes.

• per_sample_weight (bool) – Output loss for individual samples in the batch instead of mean loss.

forward(inp, target, per_sample_weights=None)

Perceptual loss forward.

Parameters

• inp (4D tensor) – Input tensor.

• target (4D tensor) – Ground truth tensor, same shape as the input.

• per_sample_weight (bool) – Output loss for individual samples in the batch instead of mean loss.

Returns

The perceptual loss.

Return type

(scalar tensor)

training = None

class imaginaire.losses.FeatureMatchingLoss(criterion='l1')

Bases: torch.nn.modules.module.Module

Compute feature matching loss

forward(fake_features, real_features)

Return the target vector for the binary cross entropy loss computation.

Parameters

• fake_features (list of lists) – Discriminator features of fake images.

• real_features (list of lists) – Discriminator features of real images.

Returns

Loss value.

Return type

(tensor)

training = None

class imaginaire.losses.GaussianKLLoss

Bases: torch.nn.modules.module.Module

Compute KL loss in VAE for Gaussian distributions

forward(mu, logvar=None)

Compute loss

Parameters

• mu (tensor) – mean

• logvar (tensor) – logarithm of variance

training = None

class imaginaire.losses.MaskedL1Loss(normalize_over_valid=False)

Bases: torch.nn.modules.module.Module

Masked L1 loss constructor.

forward(input, target, mask)

Masked L1 loss computation.

Parameters

• input (tensor) – Input tensor.

• target (tensor) – Target tensor.

• mask (tensor) – Mask to be applied to the output loss.

Returns

Loss value.

Return type

(tensor)

training = None

class imaginaire.losses.FlowLoss(cfg)

Bases: torch.nn.modules.module.Module

Flow loss constructor.

Parameters

cfg (obj) – Configuration.

compute_flow_loss(flow, warped_image, tgt_image, flow_gt, flow_conf_gt, fg_mask)

Compute losses on the generated flow map.

Parameters

• flow (tensor) – Generated flow map.

• warped_image (tensor) – Warped image using the flow map.

• tgt_image (tensor) – Target image for the warped image.

• flow_gt (tensor) – Ground truth flow map.

• flow_conf_gt (tensor) – Confidence for the ground truth flow map.

• fg_mask (tensor) – Foreground mask for the target image.

Returns

• loss_flow_L1 (tensor): L1 loss compared to ground truth flow.

• loss_flow_warp (tensor): L1 loss between the warped image and

the target image when using the flow to warp.

Return type

(dict)

compute_flow_losses(flow, warped_images, tgt_image, flow_gt, flow_conf_gt, fg_mask, tgt_label, ref_label)

Compute losses on the generated flow maps.

Parameters

• flow (tensor or list of tensors) – Generated flow maps. warped_images (tensor or list of tensors): Warped images using the flow maps.

• tgt_image (tensor) – Target image for the warped image. flow_gt (tensor or list of tensors): Ground truth flow maps.

• flow_conf_gt (tensor or list of tensors) – Confidence for the ground truth flow maps.

• fg_mask (tensor) – Foreground mask for the target image.

• tgt_label (tensor) – Target label map.

• ref_label (tensor) – Reference label map.

Returns

• loss_flow_L1 (tensor): L1 loss compared to ground truth flow.

• loss_flow_warp (tensor): L1 loss between the warped image and the target image when using the flow to warp.

• body_mask_diff (tensor): Difference between warped body part map and target body part map. Used for pose dataset only.

Return type

(dict)

compute_mask_loss(occ_mask, warped_image, tgt_image)

Compute losses on the generated occlusion mask.

Parameters

• occ_mask (tensor) – Generated occlusion mask.

• warped_image (tensor) – Warped image using the flow map.

• tgt_image (tensor) – Target image for the warped image.

Returns

Loss for the mask.

Return type

(tensor)

compute_mask_losses(occ_mask, fake_image, warped_image, tgt_label, tgt_image, fg_mask, ref_fg_mask, body_mask_diff)

Compute losses on the generated occlusion masks.

Parameters

• occ_mask (tensor or list of tensors) – Generated occlusion masks.

• fake_image (tensor) – Generated image.

• warped_image (tensor or list of tensors) – Warped images using the flow maps.

• tgt_label (tensor) – Target label map.

• tgt_image (tensor) – Target image for the warped image.

• fg_mask (tensor) – Foreground mask for the target image.

• ref_fg_mask (tensor) – Foreground mask for the reference image.

• body_mask_diff (tensor) – Difference between warped body part map

• target body part map. Used for pose dataset only. (and) –

Returns

Loss for the mask.

Return type

(tensor)

forward(data, net_G_output, current_epoch)

Compute losses on the output flow and occlusion mask.

Parameters

• data (dict) – Input data.

• net_G_output (dict) – Generator output.

• current_epoch (int) – Current training epoch number.

Returns

• loss_flow_L1 (tensor): L1 loss compared to ground truth flow.

• loss_flow_warp (tensor): L1 loss between the warped image and the target image when using the flow to warp.

• loss_mask (tensor): Loss for the occlusion mask.

Return type

(dict)

training = None

class imaginaire.losses.DictLoss(criterion='l1')

Bases: torch.nn.modules.module.Module

forward(fake, real)

Return the target vector for the l1/l2 loss computation.

Parameters

• fake (dict, list or tuple) – Discriminator features of fake images.

• real (dict, list or tuple) – Discriminator features of real images.

Returns

Loss value.

Return type

loss (tensor)

training = None

class imaginaire.losses.WeightedMSELoss(reduction='mean')

Bases: torch.nn.modules.module.Module

Compute Weighted MSE loss

forward(input, target, weight)

Return weighted MSE Loss. :param input: :type input: tensor :param target: :type target: tensor :param weight: :type weight: tensor

Returns

Loss value.

Return type

(tensor)

training = None