trw.layers.autoencoder_convolutional_variational

Module Contents

Classes

AutoencoderConvolutionalVariational

Variational convolutional autoencoder implementation
class trw.layers.autoencoder_convolutional_variational.AutoencoderConvolutionalVariational(input_shape: Union[torch.Size, List[int], Tuple[int, Ellipsis]], encoder: torch.nn.Module, decoder: torch.nn.Module, z_size: int, input_type: torch.dtype = torch.float32)

Bases: torch.nn.Module

Variational convolutional autoencoder implementation

See good reference:

https://wiseodd.github.io/techblog/2016/12/10/variational-autoencoder/

encode(self, x)
forward(self, x)
static reparameterize(training, z_mu, z_logvar)

Use the reparameterization trick: we need to generate a random normal without interrupting the gradient propagation.

We only sample during training.

static loss_function(recon_x, x, mu, logvar, recon_loss_name='BCE', kullback_leibler_weight=0.2)

Loss function generally used for a variational auto-encoder

compute:

reconstruction_loss + Kullback_Leibler_weight * Kullback–Leibler divergence((mu, logvar), gaussian(0, 1))

Parameters

  • recon_x – the reconstructed x

  • x – the input value

  • mu – the mu encoding of x

  • logvar – the logvar encoding of x

  • recon_loss_name – the name of the reconstruction loss. Must be one of BCE (binary cross-entropy) or MSE (mean squared error) or L1

  • kullback_leibler_weight – the weight factor applied on the Kullback–Leibler divergence. This is to balance the importance of the reconstruction loss and the Kullback–Leibler divergence

Returns

a 1D tensor, representing a loss value for each x

sample(self, nb_samples)

Randomly sample from the latent space to generate random samples

Parameters

nb_samples – the number of samples to generate

Notes

the image may need to be cropped or padded to mach the learnt image shape