class

VAEForImageGeneration

extendsPretrainedModel

VAEForImageGeneration(config: VAEConfig)

source edit

VAE with the ELBO training loss and prior-sample .generate().

Wraps VAEModel with the standard reconstruction + $\beta$ -weighted KL training loss and a prior-sampler convenience method. forward(x) returns a VAEOutput carrying the reconstruction, latent (concatenated across levels in HVAE mode), posterior parameters, and the three loss components (recon_loss, kl_loss, loss = recon + beta * kl). In hierarchical mode kl_loss is the sum of per-level KL terms.

generate(n_samples) samples from the prior — $z \sim \mathcal{N}(0, \mathbf{I})$ for vanilla, one $z_l \sim \mathcal{N}(0, \mathbf{I})$ per level for HVAE — and returns the decoder output. BCE reconstructions are squashed through a sigmoid before being returned.

Parameters

configVAEConfig

Hyperparameters. config.kl_weight is the

\beta

coefficient (Higgins 2017); config.recon_loss toggles between Gaussian ("mse") and Bernoulli ("bce") likelihoods.

Attributes

vaeVAEModel

Underlying VAE trunk providing encode and decode.

Notes

Reference: Kingma and Welling, "Auto-Encoding Variational Bayes", ICLR, 2014 (arXiv:1312.6114); $\beta$ -VAE in Higgins et al., ICLR, 2017; hierarchical extension in Sønderby et al., NeurIPS, 2016 (arXiv:1602.02282).

Training objective — evidence lower bound:

\mathcal{L}_{\mathrm{ELBO}} = \mathbb{E}_{q_\phi(z \mid x)}\!\big[\log p_\theta(x \mid z)\big] - \beta \cdot \mathrm{KL}\!\big( q_\phi(z \mid x) \,\big\|\, p(z) \big).

In hierarchical mode the KL term is the sum $\sum_l \mathrm{KL}(q_\phi(z_l \mid x) \| p(z_l))$ .

Examples

>>> import lucid
>>> from lucid.models.generative.vae import (
...     VAEConfig, VAEForImageGeneration,
... )
>>> cfg = VAEConfig(sample_size=32, in_channels=3, out_channels=3,
...                 latent_dim=16, down_block_channels=(16, 32, 64))
>>> model = VAEForImageGeneration(cfg).eval()
>>> x = lucid.randn((1, 3, 32, 32))
>>> out = model(x)
>>> out.sample.shape, out.loss.shape   # reconstruction + scalar loss
((1, 3, 32, 32), ())

Used by 2

Constructors

dunder

init

→None

__init__(config: VAEConfig)

source edit

Instance methods

forward

→VAEOutput

forward(x: Tensor)

source edit

generate

→GenerationOutput

generate(n_samples: int = 1, device: str = 'cpu')

source edit

Sample n_samples images from the prior N(0, I).

>>> import lucid >>> from lucid.models.generative.vae import ( ... VAEConfig, VAEForImageGeneration, ... ) >>> cfg = VAEConfig(sample_size=32, in_channels=3, out_channels=3, ... latent_dim=16, down_block_channels=(16, 32, 64)) >>> model = VAEForImageGeneration(cfg).eval() >>> x = lucid.randn((1, 3, 32, 32)) >>> out = model(x) >>> out.sample.shape, out.loss.shape # reconstruction + scalar loss ((1, 3, 32, 32), ())