VQ-VAE: Discrete Latent Representations

Module: Generative AI | Difficulty: Advanced

Vector Quantization

VQ-VAE Loss

where sg = stop gradient.

Codebook Collapse Fix (EMA Updates)

VQ-VAE-2 (Hierarchical)

import torch, torch.nn as nn

class VQVAE(nn.Module):
    def __init__(self, n_embed=512, embed_dim=64):
        super().__init__()
        self.encoder = nn.Sequential(nn.Conv2d(3,64,4,2,1),nn.ReLU(),nn.Conv2d(64,64,4,2,1),nn.ReLU())
        self.decoder = nn.Sequential(nn.ConvTranspose2d(64,64,4,2,1),nn.ReLU(),nn.ConvTranspose2d(64,3,4,2,1),nn.Tanh())
        self.vq = nn.Embedding(n_embed, embed_dim)
        self.pre_vq = nn.Conv2d(64, embed_dim, 1)
    def forward(self, x):
        ze = self.pre_vq(self.encoder(x))
        zq, indices = self.vq(ze.permute(0,2,3,1).reshape(-1,64)).reshape(ze.permute(0,2,3,1).shape).permute(0,3,1,2)
        return self.decoder(zq), ze, zq, indices