torch_mimicry.modules

Layers

Script for building specific layers needed by GAN architecture.

class ConditionalBatchNorm2d(num_features, num_classes)[source]

Conditional Batch Norm as implemented in https://github.com/pytorch/pytorch/issues/8985

num_features

Size of feature map for batch norm.

Type:int
num_classes

Determines size of embedding layer to condition BN.

Type:int
forward(x, y)[source]

Feedforwards for conditional batch norm.

Parameters:
  • x (Tensor) – Input feature map.
  • y (Tensor) – Input class labels for embedding.
Returns:

Output feature map.
Return type:

Tensor
SNConv2d(*args, default=True, **kwargs)[source]

Wrapper for applying spectral norm on conv2d layer.

SNEmbedding(*args, default=True, **kwargs)[source]

Wrapper for applying spectral norm on embedding layer.

SNLinear(*args, default=True, **kwargs)[source]

Wrapper for applying spectral norm on linear layer.

class SelfAttention(num_feat, spectral_norm=True)[source]

Self-attention layer based on version used in BigGAN code: https://github.com/ajbrock/BigGAN-PyTorch/blob/master/layers.py

forward(x)[source]

Feedforward function. Implementation differs from actual SAGAN paper, see note from BigGAN: https://github.com/ajbrock/BigGAN-PyTorch/blob/master/layers.py#L142

See official TF Implementation: https://github.com/brain-research/self-attention-gan/blob/master/non_local.py

Parameters:x (Tensor) – Input feature map.
Returns:Feature map weighed with attention map.
Return type:Tensor

Residual Blocks

Implementation of residual blocks for discriminator and generator. We follow the official SNGAN Chainer implementation as closely as possible: https://github.com/pfnet-research/sngan_projection

class DBlock(in_channels, out_channels, hidden_channels=None, downsample=False, spectral_norm=True)[source]

Residual block for discriminator.

in_channels

The channel size of input feature map.

Type:int
out_channels

The channel size of output feature map.

Type:int
hidden_channels

The channel size of intermediate feature maps.

Type:int
downsample

If True, downsamples the input feature map.

Type:bool
spectral_norm

If True, uses spectral norm for convolutional layers.

Type:bool
forward(x)[source]

Residual block feedforward function.

class DBlockOptimized(in_channels, out_channels, spectral_norm=True)[source]

Optimized residual block for discriminator. This is used as the first residual block, where there is a definite downsampling involved. Follows the official SNGAN reference implementation in chainer.

in_channels

The channel size of input feature map.

Type:int
out_channels

The channel size of output feature map.

Type:int
spectral_norm

If True, uses spectral norm for convolutional layers.

Type:bool
forward(x)[source]

Residual block feedforward function.

class GBlock(in_channels, out_channels, hidden_channels=None, upsample=False, num_classes=0, spectral_norm=False)[source]

Residual block for generator.

Uses bilinear (rather than nearest) interpolation, and align_corners set to False. This is as per how torchvision does upsampling, as seen in: https://github.com/pytorch/vision/blob/master/torchvision/models/segmentation/_utils.py

in_channels

The channel size of input feature map.

Type:int
out_channels

The channel size of output feature map.

Type:int
hidden_channels

The channel size of intermediate feature maps.

Type:int
upsample

If True, upsamples the input feature map.

Type:bool
num_classes

If more than 0, uses conditional batch norm instead.

Type:int
spectral_norm

If True, uses spectral norm for convolutional layers.

Type:bool
forward(x, y=None)[source]

Residual block feedforward function.

Losses

Loss functions definitions.

hinge_loss_dis(output_fake, output_real)[source]

Hinge loss for discriminator.

Parameters:
  • output_fake (Tensor) – Discriminator output logits for fake images.
  • output_real (Tensor) – Discriminator output logits for real images.
Returns:

A scalar tensor loss output.
Return type:

Tensor
hinge_loss_gen(output_fake)[source]

Hinge loss for generator.

Parameters:output_fake (Tensor) – Discriminator output logits for fake images.
Returns:A scalar tensor loss output.
Return type:Tensor
minimax_loss_dis(output_fake, output_real, real_label_val=1.0, fake_label_val=0.0, **kwargs)[source]

Standard minimax loss for GANs through the BCE Loss with logits fn.

Parameters:
  • output_fake (Tensor) – Discriminator output logits for fake images.
  • output_real (Tensor) – Discriminator output logits for real images.
  • real_label_val (int) – Label for real images.
  • fake_label_val (int) – Label for fake images.
  • device (torch.device) – Torch device object for sending created data.
Returns:

A scalar tensor loss output.
Return type:

Tensor
minimax_loss_gen(output_fake, real_label_val=1.0, **kwargs)[source]

Standard minimax loss for GANs through the BCE Loss with logits fn.

Parameters:
  • output (Tensor) – Discriminator output logits.
  • labels (Tensor) – Labels for computing cross entropy.
Returns:

A scalar tensor loss output.
Return type:

Tensor
ns_loss_gen(output_fake)[source]

Non-saturating loss for generator.

Parameters:output_fake (Tensor) – Discriminator output logits for fake images.
Returns:A scalar tensor loss output.
Return type:Tensor
wasserstein_loss_dis(output_real, output_fake)[source]

Computes the wasserstein loss for the discriminator.

Parameters:
  • output_real (Tensor) – Discriminator output logits for real images.
  • output_fake (Tensor) – Discriminator output logits for fake images.
Returns:

A scalar tensor loss output.
Return type:

Tensor
wasserstein_loss_gen(output_fake)[source]

Computes the wasserstein loss for generator.

Parameters:output_fake (Tensor) – Discriminator output logits for fake images.
Returns:A scalar tensor loss output.
Return type:Tensor