data
VQVAEOutput
extends
ModelOutputVQVAEOutput(sample: Tensor, latent: Tensor, indices: Tensor, perplexity: Tensor, loss: Tensor | None = None, recon_loss: Tensor | None = None, codebook_loss: Tensor | None = None, commitment_loss: Tensor | None = None)Forward output of the discrete-latent auto-encoder.
Attributes
sampleTensorReconstruction shaped
(B, C, H, W).latentTensorQuantised latent field shaped
(B, embedding_dim, H', W'), carrying the straight-through
gradient path back to the encoder.indicesTensorCodebook assignment per spatial position, shaped
(B, H', W')
with int64 dtype and values in [0, num_embeddings). This
is the discrete token field a downstream prior would model.perplexityTensorScalar over the batch's
codebook usage histogram. Ranges over
[1, num_embeddings];
a value collapsing toward 1 means the codebook has died down to a
handful of live entries.loss(Tensor or None, optional)Total objective —
recon_loss + codebook_loss + beta * commitment_loss. None on the bare
VQVAEModel, which does not build a loss.recon_loss(Tensor or None, optional)Reconstruction term alone.
codebook_loss(Tensor or None, optional) — moves the codebook
toward the encoder outputs assigned to it.
commitment_loss(Tensor or None, optional) — moves the encoder
toward the entry it selected. Reported unweighted; the
scaling is applied only inside
loss.Notes
Returned by VQVAEModel.forward (losses None) and by
VQVAEForImageGeneration.forward (losses populated).
Examples
>>> import lucid
>>> from lucid.models.generative.vqvae import (
... VQVAEConfig, VQVAEForImageGeneration,
... )
>>> cfg = VQVAEConfig(sample_size=32, hidden_channels=16,
... residual_hidden_channels=16, embedding_dim=8,
... num_embeddings=32)
>>> out = VQVAEForImageGeneration(cfg).eval()(lucid.randn((4, 3, 32, 32)))
>>> out.sample.shape, out.indices.shape
((4, 3, 32, 32), (4, 8, 8))