"""vqvae2.py — VQ-VAE-2 Hierárquico para compressão multimodal. Implementa Vector Quantized Variational Autoencoder 2 (Roy et al., 2019), uma versão hierárquica do VQ-VAE que usa dois níveis de quantização: - Top level: captura padrões semânticos globais (low-resolution) - Bottom level: captura detalhes locais (high-resolution) ============================================================================== ANÁLISE MATEMÁTICA E LÓGICA ============================================================================== ARQUITETURA: Input x ∈ R^{B×C×H×W} | v [Bottom Encoder] -> z_b ∈ R^{B×d_b×H×W} | v [Top Encoder] -> z_t ∈ R^{B×d_t×H/2×W/2} | v [Top Quantizer] -> z_t_q (codebook lookup) | v [Top Decoder] -> z_t_dec ∈ R^{B×d_b×H×W} | v [Bottom Quantizer with z_t_dec] -> z_b_q | v [Bottom Decoder] -> x_recon LOSS: L = L_recon + beta * L_commit_top + beta * L_commit_bottom + gamma * L_codebook_top + gamma * L_codebook_bottom onde: L_recon = MSE(x, x_recon) ou L1 L_commit_k = ||sg(z_k) - e_k||^2 (commit loss) L_codebook_k = ||z_k - sg(e_k)||^2 (codebook loss) sg = stop-gradient Update via EMA (Exponential Moving Average) — alternativo ao gradiente: n_k^{t+1} = gamma_ema * n_k^t + (1 - gamma_ema) * sum_k 1[assigned] m_k^{t+1} = gamma_ema * m_k^t + (1 - gamma_ema) * sum_k z_k e_k = m_k / n_k CODEBOOK: Top codebook: K_t vetores de dimensão d_t (tipicamente 512-1024 vetores) Bottom codebook: K_b vetores de dimensão d_b (tipicamente 512-1024 vetores) ============================================================================== INTEGRAÇÃO COM CNN-BiGRU ============================================================================== O VQ-VAE-2 é usado para comprimir representações multimodais (imagem, áudio) em códigos discretos que podem ser processados pelo backbone CNN-BiGRU: 1. Imagem -> VQ-VAE-2 -> códigos top + bottom (tokens discretos) 2. Códigos -> embedding lookup -> [B, T, D] 3. [B, T, D] -> CooperativeCNNBiGRU (como se fosse texto tokenizado) Isto permite que o modelo trate imagem/áudio da mesma forma que texto, unificando o pipeline multimodal. Autor: CNN-BiGRU Project """ from __future__ import annotations import logging import math from dataclasses import dataclass, field from typing import Dict, List, Optional, Tuple import torch import torch.nn as nn import torch.nn.functional as F logger = logging.getLogger(__name__) # ============================================================================ # Configuração # ============================================================================ @dataclass class VQVAE2Config: """Configuração do VQ-VAE-2 hierárquico.""" # Dimensões de entrada (assumindo features [B, C, H, W]) in_channels: int = 1 # Dimensões latentes bottom_channels: int = 64 # d_b top_channels: int = 32 # d_t # Tamanhos dos codebooks n_bottom_codes: int = 512 # K_b n_top_codes: int = 512 # K_t # Beta (commit loss weight) commitment_cost: float = 0.25 # Gamma (codebook loss weight — apenas para grad update, não EMA) codebook_cost: float = 0.0 # Usar EMA para atualizar codebook (default: True) use_ema: bool = True ema_decay: float = 0.99 # Usar residual connection (z + delta) use_residual: bool = False # Número de camadas de down/upsampling n_downsample: int = 2 # Base channels do encoder/decoder hidden_channels: int = 64 # Device device: str = "cpu" # ============================================================================ # Residual Blocks (para encoder/decoder) # ============================================================================ def _adaptive_groups(channels: int, max_groups: int = 8) -> int: """Computa número de grupos para GroupNorm que divide channels.""" if channels <= 0: return 1 g = min(max_groups, channels) while channels % g != 0 and g > 1: g -= 1 return max(1, g) class ResidualBlock(nn.Module): """Bloco residual simples com Conv2d + GroupNorm + SiLU.""" def __init__(self, channels: int, hidden_channels: Optional[int] = None): super().__init__() h = hidden_channels or channels self.conv1 = nn.Conv2d(channels, h, kernel_size=3, padding=1) self.conv2 = nn.Conv2d(h, channels, kernel_size=3, padding=1) # GroupNorm adaptativo por número de canais self.norm1 = nn.GroupNorm(_adaptive_groups(channels), channels) # norm2 aplica-se à saída de conv1 (h canais) self.norm2 = nn.GroupNorm(_adaptive_groups(h), h) def forward(self, x: torch.Tensor) -> torch.Tensor: h = F.silu(self.norm1(x)) h = self.conv1(h) h = F.silu(self.norm2(h)) h = self.conv2(h) return x + h # ============================================================================ # Encoder / Decoder (Conv2d-based) # ============================================================================ class ConvEncoder(nn.Module): """Encoder convolucional: reduz resolução espacial por fator 2^n_downsample.""" def __init__( self, in_channels: int, out_channels: int, hidden_channels: int = 64, n_downsample: int = 2, ): super().__init__() layers = [] c_in = in_channels for i in range(n_downsample): c_out = hidden_channels * (2 ** i) layers.extend([ nn.Conv2d(c_in, c_out, kernel_size=4, stride=2, padding=1), nn.GroupNorm(_adaptive_groups(c_out), c_out), nn.SiLU(), ]) c_in = c_out # Residual blocks para refinar layers.append(ResidualBlock(c_in, out_channels)) # Projeta para out_channels layers.append(nn.Conv2d(c_in, out_channels, kernel_size=1)) self.net = nn.Sequential(*layers) def forward(self, x: torch.Tensor) -> torch.Tensor: return self.net(x) class ConvDecoder(nn.Module): """Decoder convolucional: aumenta resolução espacial por fator 2^n_upsample.""" def __init__( self, in_channels: int, out_channels: int, hidden_channels: int = 64, n_upsample: int = 2, ): super().__init__() layers = [] # Primeiro projeta de in_channels para hidden_channels * 2^(n-1) c_start = hidden_channels * (2 ** (n_upsample - 1)) layers.append(nn.Conv2d(in_channels, c_start, kernel_size=1)) layers.append(ResidualBlock(c_start, c_start)) c_in = c_start for i in range(n_upsample - 1, -1, -1): c_out = hidden_channels * (2 ** max(0, i - 1)) if i > 0 else out_channels layers.extend([ nn.ConvTranspose2d(c_in, c_out, kernel_size=4, stride=2, padding=1), nn.GroupNorm(_adaptive_groups(c_out), c_out) if i > 0 else nn.Identity(), nn.SiLU() if i > 0 else nn.Identity(), ]) c_in = c_out # Camada final 1x1 para ajustar canais layers.append(nn.Conv2d(c_in, out_channels, kernel_size=1)) self.net = nn.Sequential(*layers) def forward(self, x: torch.Tensor) -> torch.Tensor: return self.net(x) # ============================================================================ # Vector Quantizer (com EMA opcional) # ============================================================================ class VectorQuantizer(nn.Module): """Quantizador por codebook com Straight-Through Estimator (STE). Args: n_codes: número de códigos no codebook (K) dim: dimensão de cada código (d) commitment_cost: peso da commitment loss use_ema: se True, atualiza codebook via EMA em vez de gradiente ema_decay: decay do EMA """ def __init__( self, n_codes: int, dim: int, commitment_cost: float = 0.25, use_ema: bool = True, ema_decay: float = 0.99, ): super().__init__() self.n_codes = n_codes self.dim = dim self.commitment_cost = commitment_cost self.use_ema = use_ema self.ema_decay = ema_decay # Codebook: [K, d] codebook = torch.randn(n_codes, dim) * 0.01 if use_ema: self.register_buffer("codebook", codebook) self.register_buffer("ema_count", torch.zeros(n_codes)) self.register_buffer("ema_weight", codebook.clone()) else: self.codebook = nn.Parameter(codebook) def forward( self, z: torch.Tensor, ) -> Tuple[torch.Tensor, torch.Tensor, Dict[str, torch.Tensor]]: """Quantiza z para o codebook mais próximo. Args: z: [B, d, H, W] ou [B, d, T] Returns: z_q: quantizado [B, d, H, W] indices: [B, H, W] índices do codebook loss_dict: dict com commitment loss, codebook loss, etc. """ # Reshape z para [N, d] onde N = B * H * W original_shape = z.shape if z.dim() == 4: B, d, H, W = z.shape z_flat = z.permute(0, 2, 3, 1).reshape(-1, d) # [B*H*W, d] elif z.dim() == 3: B, d, T = z.shape z_flat = z.permute(0, 2, 1).reshape(-1, d) # [B*T, d] else: z_flat = z # Distância euclidiana ao codebook # ||z - e||^2 = ||z||^2 - 2*z·e + ||e||^2 dist = ( z_flat.pow(2).sum(dim=-1, keepdim=True) - 2 * z_flat @ self.codebook.t() + self.codebook.pow(2).sum(dim=-1, keepdim=False).unsqueeze(0) ) # [N, K] # Encontrar código mais próximo indices = dist.argmin(dim=-1) # [N] one_hot = F.one_hot(indices, self.n_codes).float() # [N, K] # Quantizado z_q = one_hot @ self.codebook # [N, d] # Commitment loss commit_loss = F.mse_loss(z_flat, z_q.detach()) # Codebook loss (apenas se não usar EMA) if not self.use_ema: codebook_loss = F.mse_loss(z_q, z_flat.detach()) else: codebook_loss = torch.zeros(1, device=z.device) # EMA update if self.use_ema and self.training: with torch.no_grad(): # Conta quantas vezes cada código foi usado code_count = one_hot.sum(dim=0) # [K] # Soma dos z_flat para cada código code_sum = one_hot.t() @ z_flat # [K, d] # EMA self.ema_count.mul_(self.ema_decay).add_( code_count, alpha=1 - self.ema_decay ) self.ema_weight.mul_(self.ema_decay).add_( code_sum, alpha=1 - self.ema_decay ) # Atualiza codebook self.codebook.copy_(self.ema_weight / self.ema_count.unsqueeze(-1).clamp(min=1e-8)) # Straight-Through Estimator: z_q = z + (z_q - z).detach() z_q_st = z_flat + (z_q - z_flat).detach() # Reshape de volta if z.dim() == 4: z_q_st = z_q_st.view(B, H, W, d).permute(0, 3, 1, 2).contiguous() indices = indices.view(B, H, W) elif z.dim() == 3: z_q_st = z_q_st.view(B, T, d).permute(0, 2, 1).contiguous() indices = indices.view(B, T) else: z_q_st = z_q_st.view(original_shape) # Taxa de uso do codebook (para diagnóstico) with torch.no_grad(): used = (one_hot.sum(dim=0) > 0).float().sum() usage = used / self.n_codes loss_dict = { "commit_loss": commit_loss * self.commitment_cost, "codebook_loss": codebook_loss, "usage": usage.detach(), } return z_q_st, indices, loss_dict # ============================================================================ # VQ-VAE-2 Hierárquico # ============================================================================ class VQVAE2(nn.Module): """VQ-VAE-2 Hierárquico. Estrutura: x -> BottomEnc -> z_b -> TopEnc -> z_t -> TopQ -> z_t_q | v TopDec -> z_t_dec | v z_b + z_t_dec -> BottomQ -> z_b_q -> BottomDec -> x_recon """ def __init__(self, config: VQVAE2Config): super().__init__() self.config = config # Bottom encoder/decoder self.bottom_encoder = ConvEncoder( in_channels=config.in_channels, out_channels=config.bottom_channels, hidden_channels=config.hidden_channels, n_downsample=config.n_downsample, ) self.bottom_decoder = ConvDecoder( in_channels=config.bottom_channels, out_channels=config.in_channels, hidden_channels=config.hidden_channels, n_upsample=config.n_downsample, ) # Top encoder/decoder (operam sobre z_b) self.top_encoder = ConvEncoder( in_channels=config.bottom_channels, out_channels=config.top_channels, hidden_channels=config.hidden_channels, n_downsample=1, # downsample adicional por 2 ) self.top_decoder = ConvDecoder( in_channels=config.top_channels, out_channels=config.bottom_channels, hidden_channels=config.hidden_channels, n_upsample=1, ) # Quantizadores self.top_quantizer = VectorQuantizer( n_codes=config.n_top_codes, dim=config.top_channels, commitment_cost=config.commitment_cost, use_ema=config.use_ema, ema_decay=config.ema_decay, ) self.bottom_quantizer = VectorQuantizer( n_codes=config.n_bottom_codes, dim=config.bottom_channels, commitment_cost=config.commitment_cost, use_ema=config.use_ema, ema_decay=config.ema_decay, ) def forward( self, x: torch.Tensor, ) -> Dict[str, torch.Tensor]: """Forward pass completo. Args: x: [B, C, H, W] entrada (imagem, espectrograma, etc.) Returns: dict com: x_recon: [B, C, H, W] reconstrução top_indices: [B, H_t, W_t] índices top bottom_indices: [B, H_b, W_b] índices bottom loss: loss total (recon + commit + codebook) loss_dict: breakdown """ # 1. Bottom encoder z_b = self.bottom_encoder(x) # [B, d_b, H_b, W_b] # 2. Top encoder (sobre z_b) z_t = self.top_encoder(z_b) # [B, d_t, H_t, W_t] # 3. Top quantizer z_t_q, top_indices, top_loss = self.top_quantizer(z_t) # 4. Top decoder (reconstroi z_b aproximado do top) z_t_dec = self.top_decoder(z_t_q) # [B, d_b, H_b, W_b] # 5. Bottom quantizer sobre (z_b + z_t_dec) ou (z_b - z_t_dec) if self.config.use_residual: z_b_residual = z_b - z_t_dec z_b_q, bottom_indices, bottom_loss = self.bottom_quantizer(z_b_residual) z_b_full = z_b_q + z_t_dec # reconstrói else: z_b_q, bottom_indices, bottom_loss = self.bottom_quantizer(z_b) z_b_full = z_b_q + z_t_dec # combina top + bottom # 6. Bottom decoder x_recon = self.bottom_decoder(z_b_full) # [B, C, H, W] # 7. Loss recon_loss = F.mse_loss(x_recon, x) total_loss = ( recon_loss + top_loss["commit_loss"] + bottom_loss["commit_loss"] + top_loss["codebook_loss"] + bottom_loss["codebook_loss"] ) return { "x_recon": x_recon, "top_indices": top_indices, "bottom_indices": bottom_indices, "loss": total_loss, "loss_dict": { "recon_loss": recon_loss.detach(), "top_commit": top_loss["commit_loss"].detach(), "bottom_commit": bottom_loss["commit_loss"].detach(), "top_usage": top_loss["usage"], "bottom_usage": bottom_loss["usage"], }, } def encode(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: """Apenas encoder (para uso em pipeline multimodal). Returns: (top_indices, bottom_indices) — códigos discretos """ z_b = self.bottom_encoder(x) z_t = self.top_encoder(z_b) _, top_idx, _ = self.top_quantizer(z_t) z_t_dec = self.top_decoder(self.top_quantizer.codebook[top_idx].permute( 0, 3, 1, 2 ) if top_idx.dim() == 3 else self.top_quantizer.codebook[top_idx]) # Simplificado: apenas retorna top_idx e bottom_idx brutos if self.config.use_residual: z_b_resid = z_b - z_t_dec _, bottom_idx, _ = self.bottom_quantizer(z_b_resid) else: _, bottom_idx, _ = self.bottom_quantizer(z_b) return top_idx, bottom_idx def decode_codes( self, top_indices: torch.Tensor, bottom_indices: torch.Tensor, ) -> torch.Tensor: """Decodifica a partir de códigos discretos.""" # Top z_t_q = self.top_quantizer.codebook[top_indices] # [B, H_t, W_t, d_t] # Adicionar dim de canal se necessário if z_t_q.dim() == 4: z_t_q = z_t_q.permute(0, 3, 1, 2).contiguous() z_t_dec = self.top_decoder(z_t_q) # Bottom z_b_q = self.bottom_quantizer.codebook[bottom_indices] if z_b_q.dim() == 4: z_b_q = z_b_q.permute(0, 3, 1, 2).contiguous() z_b_full = z_b_q + z_t_dec # Decoder x_recon = self.bottom_decoder(z_b_full) return x_recon # ============================================================================ # Self-test # ============================================================================ def _self_test(): """Teste rápido do VQ-VAE-2.""" torch.manual_seed(42) config = VQVAE2Config( in_channels=1, bottom_channels=16, top_channels=8, n_bottom_codes=64, n_top_codes=64, n_downsample=2, hidden_channels=16, use_ema=True, commitment_cost=0.25, ) vqvae = VQVAE2(config) print(f"VQ-VAE-2 params: {sum(p.numel() for p in vqvae.parameters())}") # Forward x = torch.randn(2, 1, 16, 16) out = vqvae(x) print(f"Recon: {out['x_recon'].shape}") print(f"Top indices: {out['top_indices'].shape}") print(f"Bottom indices: {out['bottom_indices'].shape}") print(f"Loss: {out['loss'].item():.4f}") print(f"Loss dict: {out['loss_dict']}") # Test encode/decode top_idx, bot_idx = vqvae.encode(x) print(f"Encode: top {top_idx.shape}, bot {bot_idx.shape}") if __name__ == "__main__": _self_test() __all__ = [ "VQVAE2Config", "ResidualBlock", "ConvEncoder", "ConvDecoder", "VectorQuantizer", "VQVAE2", ]