API Reference

Mixin classes for VAE variants.

class vae_mixin.VAEMixin

Mixin class for basic VAE.

https://arxiv.org/abs/1312.6114

static reparameterize(mu, logvar)

Reparameterization trick.

encode(x)

Parameters:

xtorch.Tensor

The input data.

Returns:

ztorch.Tensor

The latent variables if training else mean.

decode(z)

Parameters:

ztorch.Tensor

The latent variables.

Returns:

x_rectorch.Tensor

The reconstructed data.

rec_loss(x_rec, x, **kwargs)

Compute the reconstruction loss, when the output \(p(x|z)\) is Bernoulli.

Parameters:

x_rectorch.Tensor

The reconstructed data.

xtorch.Tensor

The input data.

Returns:

bcetorch.Tensor

The binary cross entropy.

kl_loss(mu=None, logvar=None, **kwargs)

Compute the KL divergence loss, when the prior \(p(z)=\mathcal{N}(0,\mathbf{I})\) and the posterior approximation \(q(z|x)\) are Gaussian.

Parameters:

mutorch.Tensor, optional

Encoder output \(\mu\).

logvartorch.Tensor, optional

Encoder output \(\log\sigma^2\).

Returns:

kldtorch.Tensor

The KL divergence.

loss(**kwargs)

Compute the VAE objective function.

Parameters:

**kwargs

Extra arguments to loss terms.

Returns:

torch.Tensor

class vae_mixin.BetaVAEMixin(beta=1, **kwargs)

Mixin class for \(\beta\)-VAE.

https://openreview.net/forum?id=Sy2fzU9gl

Parameters:

betafloat, default 1

Regularisation coefficient \(\beta\).

loss(**kwargs)

Compute the \(\beta\)-VAE objective function.

Parameters:

**kwargs

Extra arguments to loss terms.

Returns:

torch.Tensor

class vae_mixin.InfoVAEMixin(lamb=1, alpha=0, **kwargs)

Mixin class for InfoVAE.

https://arxiv.org/abs/1706.02262

Parameters:

lambfloat, default 1

Scaling coefficient \(\lambda\).

alphafloat, default 0

Information preference \(\alpha\).

static kernel(x1, x2)

Compute the RBF kernel.

mmd_loss(z, **kwargs)

Compute the MMD loss of \(q(z)\) and \(p(z)\).

Parameters:

ztorch.Tensor

The latent variables.

Returns:

mmdtorch.Tensor

The squared maximum mean discrepancy.

loss(**kwargs)

Compute the InfoVAE objective function.

Parameters:

**kwargs

Extra arguments to loss terms.

Returns:

torch.Tensor

class vae_mixin.DIPVAEMixin(lambda_od=10, lambda_d=10, dip=1, **kwargs)

Mixin class for DIP-VAE.

https://arxiv.org/abs/1711.00848

Parameters:

lambda_odfloat, default 10

Hyperparameter \(\lambda_{od}\) for the loss on the off-diagonal entries.

lambda_dfloat, default 10

Hyperparameter \(\lambda_d\) for the loss on the diagonal entries.

dipint, 1 or 2

DIP type.

dip_regularizer(mu=None, logvar=None, **kwargs)

Compute the DIP regularization terms, when \(q(z|x)\) is Gaussian.

Parameters:

mutorch.Tensor, optional

Encoder output \(\mu\).

logvartorch.Tensor, optional

Encoder output \(\log\sigma^2\).

Returns:

off_diag_losstorch.Tensor

The off-diagonal loss.

diag_losstorch.Tensor

The diagonal loss.

loss(**kwargs)

Compute the DIP-VAE objective function.

Parameters:

**kwargs

Extra arguments to loss terms.

Returns:

torch.Tensor

class vae_mixin.BetaTCVAEMixin(alpha=1, beta=1, gamma=1, **kwargs)

Mixin class for \(\beta\)-TCVAE.

https://arxiv.org/abs/1802.04942

Parameters:

alphafloat, default 1

Weight \(\alpha\) of the index-code mutual information.

betafloat, default 1

Weight \(\beta\) of the total correlation.

gammafloat, default 1

Weight \(\gamma\) of the dimension-wise KL.

static gaussian_log_density(x, mu, logvar)

Compute the log density of a Gaussian.

decompose_kl(n, z, mu=None, logvar=None, **kwargs)

Compute the decomposed KL terms.

Parameters:

nint

Dataset size \(N\).

ztorch.Tensor

The latent variables.

mutorch.Tensor, optional

Encoder output \(\mu\).

logvartorch.Tensor, optional

Encoder output \(\log\sigma^2\).

Returns:

mi_losstorch.Tensor

The index-code mutual information.

tc_losstorch.Tensor

The total correlation.

kl_losstorch.Tensor

The dimension-wise KL.

loss(**kwargs)

Compute the \(\beta\)-TCVAE objective function.

Parameters:

**kwargs

Extra arguments to loss terms.

Returns:

torch.Tensor

class vae_mixin.VQVAEMixin(embedding_dim, num_embeddings, commitment_cost, **kwargs)

Mixin class for VQ-VAE.

https://arxiv.org/abs/1711.00937

Parameters:

embedding_dimint

The dimensionality of latent embedding vector.

num_embeddingsint

The size of the discrete latent space.

commitment_costfloat

Scalar \(\beta\) which controls the weighting of the commitment loss.

encode(x)

Parameters:

xtorch.Tensor

The input data.

Returns:

z_etorch.Tensor

The encoder outputs.

quantize(inputs)

Quantize the input vectors.

Parameters:

inputstorch.Tensor

The encoder outputs.

Returns:

z_qtorch.Tensor

The quantized vectors of the inputs.

ztorch.Tensor

The discrete latent variables.

vq_loss(z_e, z_q, **kwargs)

Compute the vector quantization loss terms.

Parameters:

z_etorch.Tensor

The encoder outputs.

z_qtorch.Tensor

The quantized vectors.

Returns:

codebook_losstorch.Tensor

commitment_losstorch.Tensor

loss(**kwargs)

Compute the VQ-VAE objective function.

Parameters:

**kwargs

Extra arguments to loss terms.

Returns:

torch.Tensor