data

VAEOutput

extendsModelOutput

VAEOutput(sample: Tensor, latent: Tensor, mu: Tensor, logvar: Tensor, loss: Tensor | None = None, recon_loss: Tensor | None = None, kl_loss: Tensor | None = None)

source edit

End-to-end VAE forward output.

Attributes

sampleTensor

Reconstructed image

\hat{x}

shaped (B, C, H, W).

latentTensor

Sampled latent

z = \mu + \sigma \cdot \varepsilon

, shape (B, latent_dim).

muTensor

Encoder mean

\mu(x)

logvarTensor

Encoder log-variance

\log \sigma^2(x)

loss(Tensor or None, optional)

Total ELBO loss (recon_loss + β · kl_loss) when targets were supplied to forward.

recon_loss(Tensor or None, optional)

Reconstruction term alone (typically BCE or MSE).

kl_loss(Tensor or None, optional)

KL divergence term

D_{\mathrm{KL}}(q(z|x) \| p(z))

alone.

Notes

Returned by VAEModel.forward. VAEForImageGeneration wraps the decoder for unconditional sampling and returns GenerationOutput from its generate method.

Examples

>>> model = create_model("vae")
>>> out = model(images)
>>> out.sample.shape, out.mu.shape
((4, 3, 64, 64), (4, 128))

Used by 2

Constructors

dunder

init

→None

__init__(sample: Tensor, latent: Tensor, mu: Tensor, logvar: Tensor, loss: Tensor | None = None, recon_loss: Tensor | None = None, kl_loss: Tensor | None = None)

source edit

data

VAEOutput

extendsModelOutput

VAEOutput(sample: Tensor, latent: Tensor, mu: Tensor, logvar: Tensor, loss: Tensor | None = None, recon_loss: Tensor | None = None, kl_loss: Tensor | None = None)

source edit

End-to-end VAE forward output.

Attributes

sampleTensor

Reconstructed image

\hat{x}

shaped (B, C, H, W).

latentTensor

Sampled latent

z = \mu + \sigma \cdot \varepsilon

, shape (B, latent_dim).

muTensor

Encoder mean

\mu(x)

logvarTensor

Encoder log-variance

\log \sigma^2(x)

loss(Tensor or None, optional)

Total ELBO loss (recon_loss + β · kl_loss) when targets were supplied to forward.

recon_loss(Tensor or None, optional)

Reconstruction term alone (typically BCE or MSE).

kl_loss(Tensor or None, optional)

KL divergence term

D_{\mathrm{KL}}(q(z|x) \| p(z))

alone.

Notes

Returned by VAEModel.forward. VAEForImageGeneration wraps the decoder for unconditional sampling and returns GenerationOutput from its generate method.

Examples

>>> model = create_model("vae")
>>> out = model(images)
>>> out.sample.shape, out.mu.shape
((4, 3, 64, 64), (4, 128))

Used by 2

Constructors

dunder

init

→None

__init__(sample: Tensor, latent: Tensor, mu: Tensor, logvar: Tensor, loss: Tensor | None = None, recon_loss: Tensor | None = None, kl_loss: Tensor | None = None)

source edit