vqvae

→VQVAEModel

vqvae(pretrained: bool = False, overrides: object = {})

source edit

Construct a VQ-VAE trunk — encoder, codebook, decoder, no loss.

Discrete-latent auto-encoder following van den Oord, Vinyals, and Kavukcuoglu, 2017, Section 4.1: two stride-2 convolutions with a 4x4 window, two pre-activation residual blocks, 256 hidden units throughout, and a 512 x 256 codebook. At the default sample_size=32 the latent grid is 8 x 8.

Use VQVAEModel.encode_indices and VQVAEModel.decode_indices for the tokeniser interface — an image in, an integer code field out, and back.

Model Size

Parameters

pretrainedbool= False

No weights are published for this family; passing True raises rather than returning a randomly initialised model.

**overridesobject= {}

Optional VQVAEConfig field overrides (e.g. sample_size=128 for the paper's ImageNet setting, num_embeddings=..., commitment_cost=...) forwarded into the underlying config.

Returns

VQVAEModel

VQ-VAE trunk configured with the paper defaults and any overrides.

Notes

Reference: van den Oord, Vinyals, and Kavukcuoglu, "Neural Discrete Representation Learning", NeurIPS, 2017 (arXiv:1711.00937).

Quantisation snaps each spatial position of the encoder output to its nearest codebook entry:

z_q(x) = e_k, \qquad k = \arg\min_j \big\| z_e(x) - e_j \big\|_2 .

Examples

>>> import lucid
>>> from lucid.models.generative.vqvae import vqvae
>>> model = vqvae().eval()
>>> x = lucid.randn((1, 3, 32, 32))
>>> out = model(x)
>>> out.sample.shape, out.indices.shape
((1, 3, 32, 32), (1, 8, 8))

Used by 1

lucid.models.generative.vqvae

>>> import lucid >>> from lucid.models.generative.vqvae import vqvae >>> model = vqvae().eval() >>> x = lucid.randn((1, 3, 32, 32)) >>> out = model(x) >>> out.sample.shape, out.indices.shape ((1, 3, 32, 32), (1, 8, 8))