Pre-processing

Load in images and apply safe data augmentation. Split to a smaller train dataset that will speed up model training. The rest are validation that don't actually touch the model but can be used to assess the quality of reconstruction from the latent bottleneck at real images on the pet manifold.

import os
os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import numpy as np
import lpips
from torch.utils.data import DataLoader, Dataset
from torchvision import transforms
from PIL import Image
import glob
import matplotlib.pyplot as plt

class DogCatDataset(Dataset):
    def __init__(self, folder_path, transform=None):
        # Collect ALL image files (both dog and cat)
        self.image_paths = glob.glob(os.path.join(folder_path, '*.jpg'))
        self.transform = transform

        # Build labels from filename: 'dog' -> 1, 'cat' -> 0
        self.labels = []
        for path in self.image_paths:
            filename = os.path.basename(path)
            if filename.startswith('dog'):
                self.labels.append(1)
            else:
                self.labels.append(0)

        # Human-readable label names (useful for plotting)
        self.label_names = {0: 'cat', 1: 'dog'}

    def __len__(self):
        return len(self.image_paths)

    def __getitem__(self, idx):
        img_path = self.image_paths[idx]
        image = Image.open(img_path).convert('RGB')
        if self.transform:
            image = self.transform(image)
        label = self.labels[idx]
        return image, label

# Define standard preprocessing
train_transform = transforms.Compose([
    transforms.Resize((72, 72)),           # slightly larger than target
    transforms.RandomCrop((64, 64)),       # random 64×64 crop
    transforms.RandomHorizontalFlip(p=0.5),
    transforms.ToTensor(),
    transforms.Lambda(lambda x: x.clamp(0, 1))
])

# Validation gets no augmentation — deterministic center crop only
val_transform = transforms.Compose([
    transforms.Resize((64, 64)),
    transforms.ToTensor(),
    transforms.Lambda(lambda x: x.clamp(0, 1))
])

dog_cat_dataset_train = DogCatDataset('dogs-vs-cats/train', transform=train_transform)
dog_cat_dataset_val   = DogCatDataset('dogs-vs-cats/train', transform=val_transform)

train_size = int(0.5 * len(dog_cat_dataset_train))
val_size   = len(dog_cat_dataset_train) - train_size

# Fix the split indices so both datasets use the same images for train/val
indices = torch.randperm(len(dog_cat_dataset_train)).tolist()
train_dataset = torch.utils.data.Subset(dog_cat_dataset_train, indices[:train_size])
val_dataset   = torch.utils.data.Subset(dog_cat_dataset_val,   indices[train_size:])

train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True,  num_workers=6)
val_loader   = DataLoader(val_dataset,   batch_size=64, shuffle=False, num_workers=6)

print(f"Number of training images:   {len(train_dataset)}")
print(f"Number of validation images: {len(val_dataset)}")

Number of training images:   12500
Number of validation images: 12500

Let's look at 100 training images to see what they look like

plt.figure(figsize=(15, 15))
for i in range(100):
    image, label = train_dataset[i]
    # Convert tensor (C, H, W) to numpy (H, W, C) for plotting
    img_np = image.permute(1, 2, 0).numpy()
    
    plt.subplot(10, 10, i + 1)
    plt.imshow(img_np)
    plt.title(dog_cat_dataset_train.label_names[label])
    plt.axis('off')

plt.tight_layout()
plt.show()

Model 1: Standard VAE w/ L1 Loss

• ResBlock Encoder w/ Stride 2 Downsampling
• ResBlock Encoder w/ Bilinear Upsampling
• Pixel likelihood conditional on latent code is assumed to follow a Laplace distribution
• KL Penalty on the latent space to force compactness

Stride 2 Downsampling vs. Max Pool

In VAEs and other generative models, it is often preferred to use Stride 2 Downsampling (via convolutional layers) rather than Max Pooling for several reasons:

1. Learnable Downsampling: Unlike Max Pooling, which is a fixed operation, a stride-step convolution allows the network to learn the most optimal way to compress spatial information and reduce dimensionality.
2. Preservation of Precise Spatial Information: Max Pooling is designed for translation invariance by discarding exact spatial location in favor of presence. In reconstruction tasks (like VAEs), knowing where a feature was is crucial for the decoder to map it back correctly.
3. Differentiability: While Max Pooling is sub-differentiable, strided convolutions provide smoother gradients throughout the training process, which helps in optimizing the bottleneck latent space.
4. Information Density: Strided convolutions help maintain a more representative summary of the input signal, whereas pooling can lead to "aliasing" or loss of fine-grained detail necessary for high-quality image synthesis.

Bilinear Interpolation vs. Transposed Convolutions

In VAEs and generative models, Bilinear Interpolation followed by a standard convolution is often preferred over Transposed Convolutions for several key reasons:

1. Elimination of Checkerboard Artifacts: Transposed convolutions frequently suffer from "checkerboard artifacts" caused by uneven overlap when the kernel size is not perfectly divisible by the stride. Bilinear upsampling provides a smooth, continuous starting point for the subsequent convolution, effectively preventing these grid-like patterns.
2. More Stable Gradient Flow: By decoupling the upsampling (interpolation) from the feature learning (convolution), the optimization process becomes more stable. The model doesn't have to learn how to both expand the spatial dimensions and refine features simultaneously within a single layer.
3. Parameter Efficiency: Bilinear interpolation is a fixed, parameter-free operation. This allows the network to focus its learnable parameters on the convolutional layers that follow, leading to a more efficient use of the model's capacity.
4. Improved Reconstruction Quality: For tasks requiring high-fidelity image synthesis, starting with a smooth bilinear estimate often leads to more natural-looking textures and edges compared to the sparse initialization inherent in transposed convolutions.

L1 vs. L2 Loss

In VAEs and other generative models, L1 Loss (Mean Absolute Error) is often preferred over L2 Loss (Mean Squared Error) for the following reasons:

1. Robustness to Outliers: L2 loss penalizes large errors quadratically, which forces the model to heavily prioritize correcting outliers. L1 loss has a constant gradient, making it more robust and less likely to be dominated by a few noisy pixels.
2. Sharper Reconstructions: L2 loss tends to produce blurry images because it averages out pixel values to minimize the squared error. L1 loss does not penalize small errors as aggressively, which encourages the model to preserve sharper edges and more distinct transitions.
3. Practical Differentiability: Although L1 is not differentiable at zero, in practice, this is rarely an issue for stochastic gradient descent. The probability of an error being exactly zero is negligible, and most deep learning frameworks (like PyTorch) provide a sub-gradient (usually 0) at that point, allowing training to proceed smoothly.

This lack of differentiability doesn't matter for VAEs where the downstream task is reconstruction. This does matter a lot for regression style problems where coefficients are going to be set to zero for certain ranges of the coefficient space leading to flat gradients!

VAE Diagram

from matplotlib.patches import FancyBboxPatch
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
 
# ════════════════════════════════════════════════════════════════
# SHARED STYLE & HELPERS
# ════════════════════════════════════════════════════════════════
 
BG_COLOR    = '#1a1a2e'
ENC_COLOR   = '#16213e'; ENC_EDGE    = '#0f3460'
DEC_COLOR   = '#1a0a2e'; DEC_EDGE    = '#6a0572'
FLAT_COLOR  = '#0a2e1a'; FLAT_EDGE   = '#0d7040'
COND_COLOR  = '#2e1a0a'; COND_EDGE   = '#d4770a'
ARROW_COLOR = '#e94560'
DIM_COLOR   = '#ffd700'
OP_COLOR    = '#a0c4ff'
TITLE_COLOR = '#ffffff'
NOTE_COLOR  = '#ff9f43'
MATH_COLOR  = '#55efc4'
GRAY        = '#aaaaaa'
GOLD_EDGE   = '#e6b800'
 
 
def draw_block(ax, x, y, w, h, title, ops, in_dim, out_dim,
               face, edge, is_flat=False):
    """Draw an architecture block with shadow, title, ops list, and I/O dims."""
    ax.add_patch(FancyBboxPatch(
        (x + 0.06, y - 0.06), w, h, boxstyle="round,pad=0.06",
        facecolor='black', edgecolor='none', alpha=0.40, zorder=1))
    ax.add_patch(FancyBboxPatch(
        (x, y), w, h, boxstyle="round,pad=0.06",
        facecolor=face, edgecolor=edge, linewidth=2.5, zorder=2))
    top = y + h
    ax.text(x + w/2, top - 0.18, f"In: {in_dim}", ha='center', va='top',
            fontsize=8.5, color=DIM_COLOR, fontweight='bold', zorder=3)
    ax.text(x + w/2, top - 0.46, title, ha='center', va='top',
            fontsize=9.5, color=TITLE_COLOR, fontweight='bold', zorder=3)
    sep = top - 0.66
    ax.plot([x + 0.10, x + w - 0.10], [sep, sep], color=edge, lw=1.0, zorder=3)
    for k, op in enumerate(ops):
        ax.text(x + w/2, sep - 0.06 - k * 0.26, op, ha='center', va='top',
                fontsize=7.8, color=OP_COLOR, zorder=3,
                style='italic' if is_flat else 'normal')
    ax.plot([x + 0.10, x + w - 0.10], [y + 0.28, y + 0.28],
            color=edge, lw=1.0, zorder=3)
    ax.text(x + w/2, y + 0.22, f"Out: {out_dim}", ha='center', va='top',
            fontsize=8.5, color=DIM_COLOR, fontweight='bold', zorder=3)
 
 
def arrow(ax, x1, y1, x2, y2, color=ARROW_COLOR, lw=2.0):
    ax.annotate('', xy=(x2, y2), xytext=(x1, y1),
                arrowprops=dict(arrowstyle='->', color=color,
                                lw=lw, mutation_scale=14),
                zorder=6)
 
 
def arrow_bend(ax, x1, y1, x2, y2, rad=0.3, color=ARROW_COLOR, lw=2.0):
    ax.annotate('', xy=(x2, y2), xytext=(x1, y1),
                arrowprops=dict(arrowstyle='->', color=color, lw=lw,
                                mutation_scale=14,
                                connectionstyle=f'arc3,rad={rad}'),
                zorder=6)
 
 
def note_box(ax, x, y, w, h, text, fontsize=10,
             face='#0f3460', edge=ARROW_COLOR, text_color=DIM_COLOR):
    """Small labeled pill for inputs/outputs."""
    ax.add_patch(FancyBboxPatch(
        (x, y), w, h, boxstyle="round,pad=0.05",
        facecolor=face, edgecolor=edge, lw=1.8, zorder=3))
    ax.text(x + w/2, y + h/2, text, ha='center', va='center',
            fontsize=fontsize, color=text_color, fontweight='bold', zorder=4)
 
 
def dashed_box(ax, x, y, w, h, edge=MATH_COLOR, lw=2.0):
    ax.add_patch(FancyBboxPatch(
        (x, y), w, h, boxstyle="round,pad=0.10",
        facecolor=BG_COLOR, edgecolor=edge,
        linewidth=lw, alpha=0.90, zorder=7, linestyle='--'))
 
 
# ════════════════════════════════════════════════════════════════
# SHARED LAYOUT (encoder / decoder positions)
# ════════════════════════════════════════════════════════════════
 
BW  = 3.00;  BH  = 2.80   # encoder / decoder block size
FBW = 2.80;  FBH = 1.95   # flat (bottleneck) block size
GAP = 0.25                 # horizontal gap between encoder blocks
 
 
def layout(enc_y, dec_y):
    """Return all x/y positions given row baselines."""
    ex = [0.2 + i * (BW + GAP) for i in range(4)]
    stack_x = ex[3] + BW + 0.55
    v_gap = 0.28
    mu_y   = enc_y + (BH - (2*FBH + v_gap))/2 + FBH + v_gap
    logv_y = enc_y + (BH - (2*FBH + v_gap))/2
    lat_x  = stack_x + FBW + 0.50
    lat_mid = logv_y + (mu_y + FBH - logv_y)/2
    lat_y  = lat_mid - FBH/2
    fcd_x  = lat_x
    fcd_y  = dec_y + (BH - FBH)/2
    dx = [ex[3], ex[2], ex[1], ex[0]]
    return dict(ex=ex, dx=dx, stack_x=stack_x,
                mu_y=mu_y, logv_y=logv_y,
                lat_x=lat_x, lat_y=lat_y,
                fcd_x=fcd_x, fcd_y=fcd_y)
 
 
# ────────────────────────────────────────────────────────────────
# Encoder & decoder specs (shared across figures)
# ────────────────────────────────────────────────────────────────
 
ENC_SPECS = [
    ("Enc Block 1",
     ["Conv 3-->64, s=1", "GN * ReLU",
      "Conv 64-->128, s=2", "GN * ReLU", "ResBlock(128)"],
     "64x64x3", "32x32x128"),
    ("Enc Block 2",
     ["Conv 128-->256, s=2", "GN * ReLU", "ResBlock(256)"],
     "32x32x128", "16x16x256"),
    ("Enc Block 3",
     ["Conv 256-->512, s=2", "GN * ReLU", "ResBlock(512)"],
     "16x16x256", "8x8x512"),
    ("Enc Block 4",
     ["Conv 512-->512, s=2", "GN * ReLU", "ResBlock(512)", "Flatten()"],
     "8x8x512", "8192"),
]
 
DEC_SPECS = [
    ("Dec Block 1",
     ["Upsample x2 (4-->8)", "Conv 512-->256",
      "GN * ReLU", "ResBlock(256)"],
     "4x4x512", "8x8x256"),
    ("Dec Block 2",
     ["Upsample x2 (8-->16)", "Conv 256-->128",
      "GN * ReLU", "ResBlock(128)"],
     "8x8x256", "16x16x128"),
    ("Dec Block 3",
     ["Upsample x2 (16-->32)", "Conv 128-->64",
      "GN * ReLU", "ResBlock(64)"],
     "16x16x128", "32x32x64"),
    ("Dec Block 4",
     ["Upsample x2 (32-->64)", "Conv 64-->32",
      "GN * ReLU", "ResBlock(32)",
      "Conv 32-->3, k=1", "Sigmoid()"],
     "32x32x64", "64x64x3"),
]
 
 
def draw_encoder(ax, L, enc_y):
    for i, (t, ops, ind, outd) in enumerate(ENC_SPECS):
        draw_block(ax, L['ex'][i], enc_y, BW, BH, t, ops, ind, outd,
                   ENC_COLOR, ENC_EDGE)
    for i in range(3):
        arrow(ax, L['ex'][i]+BW, enc_y+BH/2, L['ex'][i+1], enc_y+BH/2)
 
 
def draw_bottleneck(ax, L, enc_y):
    draw_block(ax, L['stack_x'], L['mu_y'], FBW, FBH, "fc_mu",
               ["Linear(8192 --> 256)", "-->  μ(x)  ∈ R²⁵⁶"],
               "8192", "μ  (256)", FLAT_COLOR, FLAT_EDGE, is_flat=True)
    draw_block(ax, L['stack_x'], L['logv_y'], FBW, FBH, "fc_logvar",
               ["Linear(8192 --> 256)", "-->  log σ²(x)  ∈ R²⁵⁶"],
               "8192", "log σ²  (256)", FLAT_COLOR, FLAT_EDGE, is_flat=True)
    arrow(ax, L['ex'][3]+BW, enc_y + BH*0.68, L['stack_x'], L['mu_y']+FBH/2)
    arrow(ax, L['ex'][3]+BW, enc_y + BH*0.32, L['stack_x'], L['logv_y']+FBH/2)
 
 
def draw_latent(ax, L):
    draw_block(ax, L['lat_x'], L['lat_y'], FBW, FBH,
               "Reparameterize",
               ["z = μ + σ * ε", "ε ~ N(0, I)"],
               "μ,  log σ²", "z ∈ R²⁵⁶",
               FLAT_COLOR, GOLD_EDGE, is_flat=True)
    ax.add_patch(FancyBboxPatch(
        (L['lat_x'], L['lat_y']), FBW, FBH, boxstyle="round,pad=0.06",
        facecolor='none', edgecolor=GOLD_EDGE, linewidth=3.0, zorder=5))
    arrow(ax, L['stack_x']+FBW, L['mu_y']+FBH/2,
          L['lat_x'], L['lat_y']+FBH*0.72)
    arrow(ax, L['stack_x']+FBW, L['logv_y']+FBH/2,
          L['lat_x'], L['lat_y']+FBH*0.28)
 
 
def draw_fc_decode(ax, L, in_label="z  (256)", in_size=256):
    draw_block(ax, L['fcd_x'], L['fcd_y'], FBW, FBH, "fc_decode",
               [f"Linear({in_size} --> 8192)", "view(−1, 512, 4, 4)"],
               in_label, "4x4x512", FLAT_COLOR, FLAT_EDGE, is_flat=True)
 
 
def draw_decoder(ax, L, dec_y):
    for i, (t, ops, ind, outd) in enumerate(DEC_SPECS):
        draw_block(ax, L['dx'][i], dec_y, BW, BH, t, ops, ind, outd,
                   DEC_COLOR, DEC_EDGE)
    for i in range(3):
        arrow(ax, L['dx'][i], dec_y+BH/2, L['dx'][i+1]+BW, dec_y+BH/2)
    arrow(ax, L['fcd_x'], L['fcd_y']+FBH/2, L['dx'][0]+BW, dec_y+BH/2)
 
 
# ════════════════════════════════════════════════════════════════
# FIGURE 1 — BASE VAE  +  SAMPLING NOTATION
# ════════════════════════════════════════════════════════════════
 
def figure1_base_vae():
    fig, ax = plt.subplots(1, 1, figsize=(32, 15))
    ax.set_xlim(0, 32); ax.set_ylim(0, 15)
    ax.axis('off')
    fig.patch.set_facecolor(BG_COLOR); ax.set_facecolor(BG_COLOR)
 
    enc_y, dec_y = 7.8, 2.2
    L = layout(enc_y, dec_y)
 
    # Title
    ax.text(16.0, 14.55, 'Variational Autoencoder — Architecture & Sampling',
            ha='center', va='center', fontsize=16, color=TITLE_COLOR,
            fontweight='bold', zorder=5)
 
    # Encoder
    draw_encoder(ax, L, enc_y)
    note_box(ax, L['ex'][0]-0.05, enc_y+BH+0.25, BW+0.10, 0.55,
             'Input  x ∈ R⁶⁴ˣ⁶⁴ˣ³')
    arrow(ax, L['ex'][0]+BW/2, enc_y+BH+0.25,
          L['ex'][0]+BW/2, enc_y+BH+0.02)
 
    # Bottleneck
    draw_bottleneck(ax, L, enc_y)
    draw_latent(ax, L)
 
    # Latent --> fc_decode
    arrow(ax, L['lat_x']+FBW/2, L['lat_y'],
          L['lat_x']+FBW/2, L['fcd_y']+FBH)
    draw_fc_decode(ax, L)
 
    # Decoder
    draw_decoder(ax, L, dec_y)
    note_box(ax, L['dx'][3]-0.05, dec_y-0.85, BW+0.10, 0.55,
             'Reconstruction  x_hat ∈ R⁶⁴ˣ⁶⁴ˣ³',
             face='#2e0a2e')
    arrow(ax, L['dx'][3]+BW/2, dec_y,
          L['dx'][3]+BW/2, dec_y-0.28)
 
    # ── SAMPLING NOTATION (right side) ──
    sx = L['lat_x'] + FBW + 0.40
    sy = L['lat_y'] - 0.40
    sw, sh = 5.0, 2.80
    dashed_box(ax, sx, sy, sw, sh, edge=NOTE_COLOR)
 
    ax.text(sx + sw/2, sy + sh - 0.12,
            '⚡  Generation  (no encoder needed)',
            ha='center', va='top', fontsize=11, color=NOTE_COLOR,
            fontweight='bold', zorder=8)
 
    lines = [
        ("Direct sampling:", TITLE_COLOR, True),
        ("  1.  z ~ N(0, I)           sample from prior", MATH_COLOR, False),
        ("  2.  fc_decode(z)          project to 4x4x512", MATH_COLOR, False),
        ("  3.  Decoder(.) -> x_hat   upsample to image", MATH_COLOR, False),
        ("", MATH_COLOR, False),
        ("Two-stage sampling:", TITLE_COLOR, True),
        ("  1.  z2 ~ N(0, I)          sample stage-2 prior", MATH_COLOR, False),
        ("  2.  VAE2.decode(z2) -> z1  map to stage-1 space", MATH_COLOR, False),
        ("  3.  fc_decode(z1) -> Dec -> x_hat", MATH_COLOR, False),
    ]
    for k, (txt, col, bold) in enumerate(lines):
        ax.text(sx + 0.20, sy + sh - 0.48 - k * 0.24, txt,
                ha='left', va='top', fontsize=8.5, color=col,
                fontweight='bold' if bold else 'normal',
                family='monospace', zorder=8)
 
    # ── LOSS NOTATION (below decoder) ──
    lx = L['dx'][3] + BW + 0.50
    ly = dec_y - 0.85
    lw, lh = 7.0, 1.50
    dashed_box(ax, lx, ly, lw, lh, edge=MATH_COLOR)
 
    ax.text(lx + lw/2, ly + lh - 0.10,
            'Loss  (negative ELBO)',
            ha='center', va='top', fontsize=11, color=MATH_COLOR,
            fontweight='bold', zorder=8)
    loss_lines = [
        "L  =  L1(x, x_hat)  +  lam * LPIPS(x, x_hat)  +  D_KL( Q(z|x) || N(0,I) )",
        "       |--- reconstruction ---|                   |-- regularization --|",
        "  lam  determined by gradient balancing at decoder's last layer",
    ]
    for k, line in enumerate(loss_lines):
        ax.text(lx + 0.20, ly + lh - 0.45 - k * 0.28, line,
                ha='left', va='top', fontsize=8.5, color=MATH_COLOR,
                family='monospace', zorder=8)
 
    # Section labels
    enc_mid = (L['ex'][0] + L['ex'][3] + BW) / 2
    bot_mid = (L['stack_x'] + L['lat_x'] + FBW) / 2
    dec_mid = (L['dx'][3] + L['dx'][0] + BW) / 2
    for txt, xp, yp in [("◀  ENCODER  ▶", enc_mid, 11.4),
                         ("◀  BOTTLENECK  ▶", bot_mid, 11.4),
                         ("◀  DECODER  ▶", dec_mid, 1.4)]:
        ax.text(xp, yp, txt, ha='center', va='center', fontsize=10,
                color=GRAY, fontstyle='italic', zorder=5)
 
    # Legend
    legend_items = [
        mpatches.Patch(facecolor=ENC_COLOR, edgecolor=ENC_EDGE, label='Encoder Block'),
        mpatches.Patch(facecolor=FLAT_COLOR, edgecolor=FLAT_EDGE, label='Bottleneck (Linear)'),
        mpatches.Patch(facecolor=FLAT_COLOR, edgecolor=GOLD_EDGE, label='Latent Vector z'),
        mpatches.Patch(facecolor=DEC_COLOR, edgecolor=DEC_EDGE, label='Decoder Block'),
    ]
    ax.legend(handles=legend_items, loc='lower right', fontsize=9,
              facecolor=BG_COLOR, edgecolor='white',
              labelcolor='white', framealpha=0.85)
 
    plt.tight_layout()
    plt.savefig('fig1_vae_base.png', dpi=150, bbox_inches='tight',
                facecolor=fig.get_facecolor())
    plt.show()
    print("Saved: fig1_vae_base.png")

figure1_base_vae()

Saved: fig1_vae_base.png

VAE Code

import torch
import torch.nn as nn
import torch.nn.functional as F


class ResBlock(nn.Module):
    def __init__(self, channels):
        super().__init__()
        self.block = nn.Sequential(
            nn.Conv2d(channels, channels, kernel_size=3, padding=1),
            nn.GroupNorm(32, channels),
            nn.ReLU(),
            nn.Conv2d(channels, channels, kernel_size=3, padding=1),
            nn.GroupNorm(32, channels),
        )
        self.relu = nn.ReLU()

    def forward(self, x):
        return self.relu(self.block(x) + x)


class VAE(nn.Module):
    def __init__(self, hidden_size=256):
        super(VAE, self).__init__()
        self.hidden_size = hidden_size

        self.encoder = nn.Sequential(
            nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1), # 64→64
            nn.GroupNorm(32, 32),
            nn.ReLU(),

            nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1),  # 64→32
            nn.GroupNorm(32, 64),
            nn.ReLU(),
            ResBlock(64),

            nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1),  # 32→16
            nn.GroupNorm(32, 128),
            nn.ReLU(),
            ResBlock(128),

            nn.Conv2d(128, 256, kernel_size=3, stride=2, padding=1),  # 16→8
            nn.GroupNorm(32, 256),
            nn.ReLU(),
            ResBlock(256),

            nn.Conv2d(256, 512, kernel_size=3, stride=2, padding=1),  # 8→4
            nn.GroupNorm(32, 512),
            nn.ReLU(),
            ResBlock(512),

            nn.Flatten(),
        )

        self.fc_mu = nn.Linear(512 * 4 * 4, hidden_size)
        self.fc_logvar = nn.Linear(512 * 4 * 4, hidden_size)
        self.fc_decode = nn.Linear(hidden_size, 512 * 4 * 4)

        self.decoder = nn.Sequential(

            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),  # 4→8
            nn.Conv2d(512, 256, kernel_size=3, padding=1),
            nn.GroupNorm(32, 256),
            nn.ReLU(),
            ResBlock(256),

            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),  # 8→16
            nn.Conv2d(256, 128, kernel_size=3, padding=1),
            nn.GroupNorm(32, 128),
            nn.ReLU(),
            ResBlock(128),

            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),  # 16→32
            nn.Conv2d(128, 64, kernel_size=3, padding=1),
            nn.GroupNorm(32, 64),
            nn.ReLU(),
            ResBlock(64),

            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),  # 32→64
            nn.Conv2d(64, 32, kernel_size=3, padding=1),
            nn.GroupNorm(32, 32),
            nn.ReLU(),
            ResBlock(32),

            nn.Conv2d(32, 3, kernel_size=1, stride=1),  # final projection
            nn.Sigmoid(),
        )

    def get_last_decoder_layer(self):
        """Return the weight of the final Conv2d before Sigmoid for adaptive weight calc."""
        # decoder[-2] is the Conv2d(32, 3, 1), decoder[-1] is Sigmoid
        return self.decoder[-2].weight

    def reparameterize(self, mu, logvar):
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        return mu + eps * std

    def forward(self, x):
        encoded = self.encoder(x)
        mu = self.fc_mu(encoded)
        logvar = self.fc_logvar(encoded)
        logvar = torch.clamp(logvar, min=-10, max=10)
        z = self.reparameterize(mu, logvar)
        z_projected = self.fc_decode(z).view(-1, 512, 4, 4)
        reconstruction = self.decoder(z_projected)
        return reconstruction, mu, logvar

Loss Function

Two pieces for an image $\mathbf X$ and it's reconstruction after passing through the bottleneck $\hat{\mathbf X}$. Assuming $M$ pixels in the image and $K$ dimensions at the bottleneck:

• Negative log-likelihood = $\sum \limits_{m = 1}^M | x_m - \hat{x}_m|$

• KL Divergence w/ $K$ dim standard normal prior = $\sum \limits_{k = 1}^K -\frac{1}{2} \left[1 + \log \sigma^2_k - \mu_k^2 - \sigma^2_k \right]$

ELBO derivation says just add 'em together and minimize!

def generator_loss(reconstruction, x, mu, logvar, beta=1.0, free_bits=0.1):

    #n_pixels = x.shape[1] * x.shape[2] * x.shape[3]

    # L1 at ELBO scale (sum over pixels, mean over batch)
    # This is the natural negative likelihood loss!  Must be this scale for the KL
    # Term to Balance

    l1_loss = F.l1_loss(reconstruction, x, reduction='sum') / x.shape[0]

    nll_loss = l1_loss

    # KL divergence (sum over latent dims, mean over batch)
    # Normal Posterior vs. Normal Prior
    # Independent, so final solution at log scale is sum
    KLD_per_dim = -0.5 * (1 + logvar - mu.pow(2) - logvar.exp())
    # Any KL loss by dim less than free_bits is coded as free_bits
    # Removes incentive at the margins to exactly match prior
    KLD = torch.sum(torch.clamp(KLD_per_dim, min=free_bits)) / x.shape[0]

    total = nll_loss + (KLD * beta)

    return total, nll_loss.item(), KLD.item()

Model Setup

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
hidden_size = 256

model = VAE(hidden_size=hidden_size).to(device)

print(model)

VAE(
  (encoder): Sequential(
    (0): Conv2d(3, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): GroupNorm(32, 32, eps=1e-05, affine=True)
    (2): ReLU()
    (3): Conv2d(32, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
    (4): GroupNorm(32, 64, eps=1e-05, affine=True)
    (5): ReLU()
    (6): ResBlock(
      (block): Sequential(
        (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (1): GroupNorm(32, 64, eps=1e-05, affine=True)
        (2): ReLU()
        (3): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (4): GroupNorm(32, 64, eps=1e-05, affine=True)
      )
      (relu): ReLU()
    )
    (7): Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
    (8): GroupNorm(32, 128, eps=1e-05, affine=True)
    (9): ReLU()
    (10): ResBlock(
      (block): Sequential(
        (0): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (1): GroupNorm(32, 128, eps=1e-05, affine=True)
        (2): ReLU()
        (3): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (4): GroupNorm(32, 128, eps=1e-05, affine=True)
      )
      (relu): ReLU()
    )
    (11): Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
    (12): GroupNorm(32, 256, eps=1e-05, affine=True)
    (13): ReLU()
    (14): ResBlock(
      (block): Sequential(
        (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (1): GroupNorm(32, 256, eps=1e-05, affine=True)
        (2): ReLU()
        (3): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (4): GroupNorm(32, 256, eps=1e-05, affine=True)
      )
      (relu): ReLU()
    )
    (15): Conv2d(256, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
    (16): GroupNorm(32, 512, eps=1e-05, affine=True)
    (17): ReLU()
    (18): ResBlock(
      (block): Sequential(
        (0): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (1): GroupNorm(32, 512, eps=1e-05, affine=True)
        (2): ReLU()
        (3): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (4): GroupNorm(32, 512, eps=1e-05, affine=True)
      )
      (relu): ReLU()
    )
    (19): Flatten(start_dim=1, end_dim=-1)
  )
  (fc_mu): Linear(in_features=8192, out_features=256, bias=True)
  (fc_logvar): Linear(in_features=8192, out_features=256, bias=True)
  (fc_decode): Linear(in_features=256, out_features=8192, bias=True)
  (decoder): Sequential(
    (0): Upsample(scale_factor=2.0, mode='bilinear')
    (1): Conv2d(512, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (2): GroupNorm(32, 256, eps=1e-05, affine=True)
    (3): ReLU()
    (4): ResBlock(
      (block): Sequential(
        (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (1): GroupNorm(32, 256, eps=1e-05, affine=True)
        (2): ReLU()
        (3): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (4): GroupNorm(32, 256, eps=1e-05, affine=True)
      )
      (relu): ReLU()
    )
    (5): Upsample(scale_factor=2.0, mode='bilinear')
    (6): Conv2d(256, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (7): GroupNorm(32, 128, eps=1e-05, affine=True)
    (8): ReLU()
    (9): ResBlock(
      (block): Sequential(
        (0): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (1): GroupNorm(32, 128, eps=1e-05, affine=True)
        (2): ReLU()
        (3): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (4): GroupNorm(32, 128, eps=1e-05, affine=True)
      )
      (relu): ReLU()
    )
    (10): Upsample(scale_factor=2.0, mode='bilinear')
    (11): Conv2d(128, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (12): GroupNorm(32, 64, eps=1e-05, affine=True)
    (13): ReLU()
    (14): ResBlock(
      (block): Sequential(
        (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (1): GroupNorm(32, 64, eps=1e-05, affine=True)
        (2): ReLU()
        (3): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (4): GroupNorm(32, 64, eps=1e-05, affine=True)
      )
      (relu): ReLU()
    )
    (15): Upsample(scale_factor=2.0, mode='bilinear')
    (16): Conv2d(64, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (17): GroupNorm(32, 32, eps=1e-05, affine=True)
    (18): ReLU()
    (19): ResBlock(
      (block): Sequential(
        (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (1): GroupNorm(32, 32, eps=1e-05, affine=True)
        (2): ReLU()
        (3): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (4): GroupNorm(32, 32, eps=1e-05, affine=True)
      )
      (relu): ReLU()
    )
    (20): Conv2d(32, 3, kernel_size=(1, 1), stride=(1, 1))
    (21): Sigmoid()
  )
)

AdamW optimizer. Learning rate is set using a cosine annealing procedure with warm restarts every 50x iterations

from torch import optim

# Optimization
# AdamW with default weight decay
optimizer = optim.AdamW(model.parameters(), lr=1e-3)
#Long burn
num_epochs = 500


# Cosine w/ Warm Restarts
scheduler = optim.lr_scheduler.CosineAnnealingWarmRestarts(
    optimizer,
    T_0=50,          # First cycle is 50 epochs
    T_mult=2,        # Double the length of each subsequent cycle
    eta_min=1e-6     # Let it decay very close to zero
)

Slowly turn on the KL annealing to prevent immediate collapse. 5 iterations with no regularization, 10 slowly ramping from 0 KL penalty to the correct KL penalty!

## KL Annealing

import numpy as np

def get_beta(epoch, warmup_start=10, warmup_end=20, beta_max=1.0):
    if epoch <= warmup_start:
        return 0.0
    elif epoch <= warmup_end:
        t = epoch - warmup_start
        warmup_steps = warmup_end - warmup_start
        return beta_max * 0.5 * (1 - np.cos(np.pi * t / warmup_steps))
    else:
        return beta_max

A function to allow us to visualize evolution of image reconstructions and to show the decoded form of generated images from the latent space:

def visualize_results(model, epoch, temperature=0.7):
    model.eval()
    with torch.no_grad():
        train_batch, _ = next(iter(train_loader))
        train_batch = train_batch[:5].to(device)
        train_recon, _, _ = model(train_batch)

        val_batch, _ = next(iter(val_loader))
        val_batch = val_batch[:5].to(device)
        val_recon, _, _ = model(val_batch)

        z_random = torch.randn(10, hidden_size).to(device) * temperature
        z_projected = model.fc_decode(z_random).view(-1, 512, 4, 4)
        generated = model.decoder(z_projected)

        fig, axes = plt.subplots(3, 10, figsize=(20, 6))

        for i in range(5):
            axes[0, i * 2].imshow(train_batch[i].cpu().permute(1, 2, 0))
            axes[0, i * 2].set_title("Train Ori")
            axes[0, i * 2 + 1].imshow(train_recon[i].cpu().permute(1, 2, 0))
            axes[0, i * 2 + 1].set_title("Train Recon")

            axes[1, i * 2].imshow(val_batch[i].cpu().permute(1, 2, 0))
            axes[1, i * 2].set_title("Val Ori")
            axes[1, i * 2 + 1].imshow(val_recon[i].cpu().permute(1, 2, 0))
            axes[1, i * 2 + 1].set_title("Val Recon")

        for i in range(10):
            axes[2, i].imshow(generated[i].cpu().permute(1, 2, 0))
            axes[2, i].set_title(f"Gen {i + 1}")

        for ax in axes.flatten():
            ax.axis("off")

        plt.suptitle(f"Epoch {epoch}")
        plt.tight_layout()
        plt.show()

Now, our training loop (with some additional niceties that will make for better visualization):

from tqdm import tqdm
import numpy as np
import torch
import plotly.graph_objects as go
from plotly.subplots import make_subplots
import ipywidgets as widgets
from IPython.display import display
import base64, io
from PIL import Image as PILImage

# =============================================================================
# DASHBOARD SETUP  — Plotly + ipywidgets (survives nbconvert → HTML)
# =============================================================================

history = {'epoch': [], 'total': [], 'nll': [], 'kld': []}
snapshot_store = {}   # epoch → list of 30 PIL images

# ── Loss figure (Plotly FigureWidget — updates in-place) ──────────
loss_fig = make_subplots(
    rows=1, cols=3,
    subplot_titles=('Total ELBO Loss', 'NLL  (Reconstruction)', 'KL Divergence'),
)
loss_fig.add_trace(go.Scatter(x=[], y=[], mode='lines',
                               line=dict(color='#e94560', width=2),
                               name='Total'), row=1, col=1)
loss_fig.add_trace(go.Scatter(x=[], y=[], mode='lines',
                               line=dict(color='#a0c4ff', width=2),
                               name='NLL'),   row=1, col=2)
loss_fig.add_trace(go.Scatter(x=[], y=[], mode='lines',
                               line=dict(color='#ffd700', width=2),
                               name='KLD'),   row=1, col=3)
loss_fig.update_layout(
    height=320, paper_bgcolor='#1a1a2e', plot_bgcolor='#0d0d1a',
    font=dict(color='white'), showlegend=False,
    margin=dict(t=40, b=30, l=40, r=20),
)
for i in range(1, 4):
    loss_fig.update_xaxes(gridcolor='#222244', title_text='Epoch',
                           row=1, col=i)
    loss_fig.update_yaxes(gridcolor='#222244', row=1, col=i)

loss_widget = go.FigureWidget(loss_fig)


def _update_loss_plots():
    ep  = history['epoch']
    with loss_widget.batch_update():
        loss_widget.data[0].x = ep;  loss_widget.data[0].y = history['total']
        loss_widget.data[1].x = ep;  loss_widget.data[1].y = history['nll']
        loss_widget.data[2].x = ep;  loss_widget.data[2].y = history['kld']


# ── Snapshot viewer (ipywidgets — buttons + image grid) ───────────
def _tensor_to_b64(img_np):
    """Convert H×W×3 float32 numpy array → base64 PNG string."""
    arr = (np.clip(img_np, 0, 1) * 255).astype(np.uint8)
    buf = io.BytesIO()
    PILImage.fromarray(arr).save(buf, format='PNG')
    return 'data:image/png;base64,' + base64.b64encode(buf.getvalue()).decode()

ROW_LABELS = ['Train  Orig / Recon', 'Val  Orig / Recon', 'Generated']

# 30 image widgets arranged in 3 rows × 10 cols
img_widgets = [[widgets.Image(format='png', width=80, height=80)
                for _ in range(10)] for _ in range(3)]

row_boxes = []
for r in range(3):
    label = widgets.Label(
        value=ROW_LABELS[r],
        layout=widgets.Layout(width='130px', display='flex',
                              align_items='center'))
    row_box = widgets.HBox(
        [label] + img_widgets[r],
        layout=widgets.Layout(align_items='center'))
    row_boxes.append(row_box)

snap_label = widgets.Label(
    value='Snapshot — no data yet',
    layout=widgets.Layout(margin='4px 0 4px 0'))
snap_label.style.font_size = '13px'

snap_idx = [0]

btn_prev = widgets.Button(description='◀  Prev',
                          button_style='',
                          layout=widgets.Layout(width='100px'))
btn_next = widgets.Button(description='Next  ▶',
                          button_style='',
                          layout=widgets.Layout(width='100px'))
for b in (btn_prev, btn_next):
    b.style.button_color = '#16213e'
    b.style.text_color   = 'white'

epoch_dropdown = widgets.Dropdown(
    options=[],
    description='Epoch:',
    layout=widgets.Layout(width='150px'))


def _render_snapshot(epoch):
    imgs = snapshot_store.get(epoch)
    if imgs is None:
        return
    for r in range(3):
        for c in range(10):
            b64 = _tensor_to_b64(imgs[r * 10 + c])
            # FigureWidget images need raw bytes
            raw = base64.b64decode(b64.split(',')[1])
            img_widgets[r][c].value = raw
    snap_label.value = f'Snapshot — Epoch {epoch}'


def _on_prev(_b):
    keys = sorted(snapshot_store.keys())
    if not keys: return
    snap_idx[0] = max(0, snap_idx[0] - 1)
    epoch_dropdown.value = keys[snap_idx[0]]


def _on_next(_b):
    keys = sorted(snapshot_store.keys())
    if not keys: return
    snap_idx[0] = min(len(keys) - 1, snap_idx[0] + 1)
    epoch_dropdown.value = keys[snap_idx[0]]


def _on_dropdown_change(change):
    if change['name'] == 'value' and change['new'] is not None:
        keys = sorted(snapshot_store.keys())
        snap_idx[0] = keys.index(change['new'])
        _render_snapshot(change['new'])


btn_prev.on_click(_on_prev)
btn_next.on_click(_on_next)
epoch_dropdown.observe(_on_dropdown_change)

nav_bar      = widgets.HBox([btn_prev, epoch_dropdown, btn_next])
snap_viewer  = widgets.VBox([snap_label, nav_bar] + row_boxes)
dashboard    = widgets.VBox([loss_widget, snap_viewer])
display(dashboard)


def _capture_snapshot(epoch):
    model.eval()
    imgs = []
    with torch.no_grad():
        tb, _ = next(iter(train_loader))
        tb = tb[:5].to(device);  tr, _, _ = model(tb)
        for i in range(5):
            imgs.append(tb[i].cpu().permute(1,2,0).numpy())
            imgs.append(tr[i].cpu().permute(1,2,0).numpy())

        vb, _ = next(iter(val_loader))
        vb = vb[:5].to(device);  vr, _, _ = model(vb)
        for i in range(5):
            imgs.append(vb[i].cpu().permute(1,2,0).numpy())
            imgs.append(vr[i].cpu().permute(1,2,0).numpy())

        z_rand = torch.randn(10, hidden_size).to(device) * 0.7
        gen = model.decoder(model.fc_decode(z_rand).view(-1, 512, 4, 4))
        for i in range(10):
            imgs.append(gen[i].cpu().permute(1,2,0).numpy())

    snapshot_store[epoch] = imgs
    # Update dropdown options
    epoch_dropdown.options = sorted(snapshot_store.keys())
    snap_idx[0] = len(snapshot_store) - 1
    epoch_dropdown.value = epoch
    model.train()

# =============================================================================
# TRAINING LOOP  — plain VAE, single optimizer
# =============================================================================

best_avg_loss = float('inf')
global_step   = 0

for epoch in range(1, num_epochs + 1):
    model.train()
    train_loss = 0.0
    total_nll  = 0.0
    total_kld  = 0.0

    beta = get_beta(epoch)
    #beta = 1
    pbar = tqdm(train_loader,
                desc=f"Epoch {epoch:3d}/{num_epochs}", leave=False)

    for batch_idx, (data, _) in enumerate(pbar):
        data = data.to(device)

        optimizer.zero_grad()
        recon_batch, mu, logvar = model(data)

        loss, nll, kld = generator_loss(
            reconstruction = recon_batch,
            x              = data,
            mu             = mu,
            logvar         = logvar,
            beta           = beta,
            free_bits      = 0.1,
        )

        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=10.0)
        optimizer.step()

        train_loss += loss.item()
        total_nll  += nll
        total_kld  += kld
        global_step += 1

        pbar.set_postfix({
            'Loss': f'{train_loss / (batch_idx+1):.1f}',
            'NLL':  f'{total_nll  / (batch_idx+1):.1f}',
            'KLD':  f'{total_kld  / (batch_idx+1):.1f}',
            'β':    f'{beta:.3f}',
        })

    # ── Epoch averages ────────────────────────────────────────────
    n          = len(train_loader)
    avg_loss   = train_loss / n
    avg_nll    = total_nll  / n
    avg_kld    = total_kld  / n
    current_lr = optimizer.param_groups[0]['lr']

    # ── Checkpoint on best training ELBO ─────────────────────────
    if avg_loss < best_avg_loss:
        best_avg_loss = avg_loss
        torch.save({
            "epoch":            epoch,
            "global_step":      global_step,
            "model_state_dict": model.state_dict(),
            "opt_state_dict":   optimizer.state_dict(),
            "best_avg_loss":    best_avg_loss,
            "hidden_size":      model.hidden_size,
        }, "vae_best.pth")

    scheduler.step()

    # ── Live dashboard update ─────────────────────────────────────
    if epoch >= 20:
        history['epoch'].append(epoch)
        history['total'].append(avg_loss)
        history['nll']  .append(avg_nll)
        history['kld']  .append(avg_kld)
        _update_loss_plots()

    # ── Snapshot every 10 epochs ──────────────────────────────────
    if epoch % 10 == 0:
        _capture_snapshot(epoch)

VBox(children=(FigureWidget({
    'data': [{'line': {'color': '#e94560', 'width': 2},
              'mode': 'l…

Model 2: VAE + Perceptual Loss

Our reconstructions and generated images seem right - they look like dogs and cats. But, it's hard to actually say if they are or aren't because the images are incredibly blurry!

This is a problem with our choice of loss function - through the ELBO only requires us to minimize NLL + KL, minimizing the NLL leads to behavior that really penalizes shifts from the original position. As an example, we want to construct our latent space in such a way that cat in returns a cat out. Since our loss function is per-pixel loss, a cat outline shifted just one pixel to the right drastically increases the loss! To the human eye, though, a cat that is one pixel wider than it should be is unimportant - a cat is a cat is a cat.

As with CNNs and transitioning to U-Nets, images have two properties - what and where. Our VAE as is does a decent job of handling both, but the loss structure leads the VAE to prioritize where more than what!

An easy solution is to try to measure what and add a little where. This can be accomplished by combining the per-pixel loss with perceptual loss.

Perceptual loss leverages the convolutional backbone of a CNN trained for classification and says that two identical images should have the same feature maps throughout the convolutional backbone. For a 64 x 64 x 3 image and a 64 x 64 x 3 reconstruction, we should see that the feature map for this image given in the first layer of a CNN backbone (32 height x 32 width x 64 channels, for example) should be the same for both images if they are identical.

As CNN backbones are constructed to induce hierarchical feature arrangements - local patterns early and global patterns late - the difference between the feature maps for a true image and its reconstruction at each layer contains different levels of information for what and where. Something like the absolute value of the differences of these feature maps would provide more information about what is actually different between the images, regardless of where.

Let PL denote a measure of perceptual loss. Then, we can define a loss function to minimize as:

$$ \text{PL} + \text{NLL} + \text{KLD} $$

The most common way to measure perceptual loss is using LPIPS (Learned Perceptual Image Patch Similarity; Zhang et al. 2018). LPIPS uses a pretrained VGG ImageNet backbone and passes both an image and its reconstruction through and captures all intermediate representations. Then,

1. Channel-wise normalization: At each layer, the feature activations are normalized to unit length along the channel dimension. This ensures that all channels contribute equally regardless of their raw activation magnitudes — without this step, a channel with activations in the range [0, 100] would dominate the loss compared to a channel in [0, 1], regardless of which is more perceptually meaningful.

2. Learned per-channel weighting: Each channel is multiplied by a scalar weight that was trained on human perceptual judgments. Zhang et al. collected a dataset (BAPPS) where humans were shown image triplets and asked "which of these two patches is more similar to the reference?" The learned weights encode which VGG channels humans actually care about — amplifying channels that detect perceptually important features (edges, textures, object parts) and suppressing channels that detect things humans are insensitive to.

3. Squared difference, averaged spatially: At each layer $l$, the squared difference of the weighted, normalized features is averaged over spatial positions:

$$d_l = \frac{1}{H_l W_l} \sum_{h,w} \left\| w_l \odot \left( \hat{x}_l^{h,w} - \hat{y}_l^{h,w} \right) \right\|_2^2$$

4. Sum across layers: The final LPIPS score sums across all extracted layers:

$$\text{LPIPS}(x, \hat{x}) = \sum_l d_l$$

The critical insight is that VGG features capture spatial relationships between pixels. Early layers detect edges and textures. Later layers detect object parts and compositions. Two images can have low LPIPS loss even if they differ pixel-by-pixel — as long as they share the same visual structure. A cat shifted one pixel to the right has nearly identical VGG features to the original cat, so the perceptual loss is close to zero even though the per-pixel loss is large.

This frees the decoder from the tyranny of pixel-perfect alignment. Instead of hedging its bets and producing a blurry average, it can commit to sharp, coherent outputs — because the loss function now rewards perceptual similarity rather than pixel identity.

The tradeoff: LPIPS has no probabilistic interpretation. The ELBO derivation gives us NLL + KLD, and bolting LPIPS onto that is theoretically unprincipled — we're adding a term that doesn't correspond to any likelihood model. And this can create problems! If the perceptual loss on its natural scale is much bigger than the NLL, then the reconstructions will prioritize structure over location. If the opposite is true, then the reconstructions will prioritize location over structure. And if these two things are really big or small and not appropriately weighted, then the KL divergence could completely dominate the optimization procedure (posterior collapse to the prior) or get completely lost (shift to deterministic autoencoders). The best solution is to find a principled way to balance these terms through gradient balancing.

Gradient Balancing

The core insight is that loss magnitudes don't matter — gradient magnitudes do. The optimizer doesn't see loss values directly. It sees gradients. Two loss terms can have identical values but wildly different gradient norms at the decoder's weights, meaning one term dominates the parameter update while the other has no practical influence.

So instead of comparing loss values and guessing a weight, we compare gradient norms at a specific point in the network — typically the decoder's last layer, since that's where all reconstruction signals converge before producing the output image.

Given two loss terms $\mathcal{L}_A$ and $\mathcal{L}_B$ that we want to balance, we compute:

$$\lambda = \frac{\left\| \nabla_{w} \mathcal{L}_A \right\|}{\left\| \nabla_{w} \mathcal{L}_B \right\| + \epsilon}$$

where $w$ are the weights of the decoder's last layer and $\epsilon$ is a small constant for numerical stability. Then we use $\lambda$ to scale $\mathcal{L}_B$:

$$\mathcal{L}_{\text{total}} = \mathcal{L}_A + \lambda \cdot \mathcal{L}_B + \mathcal{L}_{\text{KLD}}$$

When $\lambda$ is large, it means $\mathcal{L}_A$ produces much larger gradients than $\mathcal{L}_B$, so $\mathcal{L}_B$ needs to be amplified to compete. When $\lambda$ is small, $\mathcal{L}_B$ is already strong enough relative to $\mathcal{L}_A$.

Why this works: $\lambda$ is recomputed at every training step, so it adapts automatically as the loss landscape changes during training. Early in training when reconstructions are poor, the gradient norms might favor one term heavily. Later when reconstructions improve, the balance shifts. A fixed weight can't track this — gradient balancing can.

Why the last decoder layer? This is the final point where the network can influence the output image. If a loss term has zero gradient at this layer, it has zero influence on what pixels are produced, regardless of how large the loss value is. Measuring gradients here tells us exactly how much each term is actually steering the decoder's output.

In PyTorch, the computation uses torch.autograd.grad to peek at the gradients without performing a full backward pass:

def compute_adaptive_weight(loss_a, loss_b, last_layer_weight):
    grads_a = torch.autograd.grad(loss_a, last_layer_weight, retain_graph=True)[0]
    grads_b = torch.autograd.grad(loss_b, last_layer_weight, retain_graph=True)[0]
    weight = torch.norm(grads_a) / (torch.norm(grads_b) + 1e-4)
    weight = torch.clamp(weight, 0.0, 1e4).detach()
    return weight

Note the .detach() — $\lambda$ is treated as a constant for the backward pass. We don't want to differentiate through the balancing computation itself; we just want it to set the scale.

For our VAE with perceptual loss, we set $\mathcal{L}_A$ = NLL (the ELBO-grounded term we trust) and $\mathcal{L}_B$ = LPIPS (the term with no natural scale). This makes NLL the anchor and dynamically scales LPIPS to have proportional gradient influence at the decoder's output layer. The KLD term is left unscaled because it operates on the encoder's output (mu and logvar), not the decoder — its gradients flow through a completely different path and are already at the correct ELBO scale relative to NLL.

Let's start building this improved VAE! We'll use the exact same VAE architecture as before. The only thing that changes is the Loss computation.

First, we need to create the LPIPS module that will allow us to compute the perceptual loss:

import lpips

class LPIPSLoss(nn.Module):
    """LPIPS perceptual loss. Expects [0, 1] inputs, converts to [-1, 1] internally."""
    def __init__(self):
        super().__init__()
        self.lpips = lpips.LPIPS(net='vgg')
        for param in self.parameters():
            param.requires_grad = False

    def forward(self, reconstruction, target):
        return self.lpips(reconstruction * 2 - 1, target * 2 - 1).mean()

Next, we need to set up the adaptive weight function:

def compute_adaptive_weight(loss_a, loss_b, last_layer_weight):
    grads_a = torch.autograd.grad(loss_a, last_layer_weight, retain_graph=True)[0]
    grads_b = torch.autograd.grad(loss_b, last_layer_weight, retain_graph=True)[0]
    weight = torch.norm(grads_a) / (torch.norm(grads_b) + 1e-4)
    weight = torch.clamp(weight, 0.0, 1e4).detach()
    return weight

Then, we need to set up our generator loss function:

def perceptual_generator_loss(reconstruction, x, mu, logvar, perceptual_loss_fn, last_layer_weight, beta=1.0, free_bits=0.1):
    # NLL Loss (L1 at ELBO scale)
    nll_loss = F.l1_loss(reconstruction, x, reduction='sum') / x.shape[0]

    # Perceptual Loss via LPIPS
    p_loss = perceptual_loss_fn(reconstruction, x)

    # KL Divergence
    KLD_per_dim = -0.5 * (1 + logvar - mu.pow(2) - logvar.exp())
    KLD = torch.sum(torch.clamp(KLD_per_dim, min=free_bits)) / x.shape[0]

    # Gradient balancing: compute adaptive weight to scale perceptual loss to NLL scale
    d_weight = compute_adaptive_weight(nll_loss, p_loss, last_layer_weight)

    wp_loss = d_weight * p_loss

    # Total ELBO with balanced perceptual term
    total = nll_loss + (wp_loss) + (beta * KLD)

    return total, nll_loss.item(), KLD.item(), wp_loss.item(), d_weight.item()

As before, we need to instantiate our models (VAE and LPIPS) and start up a new optimizer.

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
hidden_size = 256

model = VAE(hidden_size=hidden_size).to(device)

print(model)

VAE(
  (encoder): Sequential(
    (0): Conv2d(3, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): GroupNorm(32, 32, eps=1e-05, affine=True)
    (2): ReLU()
    (3): Conv2d(32, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
    (4): GroupNorm(32, 64, eps=1e-05, affine=True)
    (5): ReLU()
    (6): ResBlock(
      (block): Sequential(
        (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (1): GroupNorm(32, 64, eps=1e-05, affine=True)
        (2): ReLU()
        (3): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (4): GroupNorm(32, 64, eps=1e-05, affine=True)
      )
      (relu): ReLU()
    )
    (7): Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
    (8): GroupNorm(32, 128, eps=1e-05, affine=True)
    (9): ReLU()
    (10): ResBlock(
      (block): Sequential(
        (0): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (1): GroupNorm(32, 128, eps=1e-05, affine=True)
        (2): ReLU()
        (3): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (4): GroupNorm(32, 128, eps=1e-05, affine=True)
      )
      (relu): ReLU()
    )
    (11): Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
    (12): GroupNorm(32, 256, eps=1e-05, affine=True)
    (13): ReLU()
    (14): ResBlock(
      (block): Sequential(
        (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (1): GroupNorm(32, 256, eps=1e-05, affine=True)
        (2): ReLU()
        (3): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (4): GroupNorm(32, 256, eps=1e-05, affine=True)
      )
      (relu): ReLU()
    )
    (15): Conv2d(256, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
    (16): GroupNorm(32, 512, eps=1e-05, affine=True)
    (17): ReLU()
    (18): ResBlock(
      (block): Sequential(
        (0): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (1): GroupNorm(32, 512, eps=1e-05, affine=True)
        (2): ReLU()
        (3): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (4): GroupNorm(32, 512, eps=1e-05, affine=True)
      )
      (relu): ReLU()
    )
    (19): Flatten(start_dim=1, end_dim=-1)
  )
  (fc_mu): Linear(in_features=8192, out_features=256, bias=True)
  (fc_logvar): Linear(in_features=8192, out_features=256, bias=True)
  (fc_decode): Linear(in_features=256, out_features=8192, bias=True)
  (decoder): Sequential(
    (0): Upsample(scale_factor=2.0, mode='bilinear')
    (1): Conv2d(512, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (2): GroupNorm(32, 256, eps=1e-05, affine=True)
    (3): ReLU()
    (4): ResBlock(
      (block): Sequential(
        (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (1): GroupNorm(32, 256, eps=1e-05, affine=True)
        (2): ReLU()
        (3): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (4): GroupNorm(32, 256, eps=1e-05, affine=True)
      )
      (relu): ReLU()
    )
    (5): Upsample(scale_factor=2.0, mode='bilinear')
    (6): Conv2d(256, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (7): GroupNorm(32, 128, eps=1e-05, affine=True)
    (8): ReLU()
    (9): ResBlock(
      (block): Sequential(
        (0): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (1): GroupNorm(32, 128, eps=1e-05, affine=True)
        (2): ReLU()
        (3): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (4): GroupNorm(32, 128, eps=1e-05, affine=True)
      )
      (relu): ReLU()
    )
    (10): Upsample(scale_factor=2.0, mode='bilinear')
    (11): Conv2d(128, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (12): GroupNorm(32, 64, eps=1e-05, affine=True)
    (13): ReLU()
    (14): ResBlock(
      (block): Sequential(
        (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (1): GroupNorm(32, 64, eps=1e-05, affine=True)
        (2): ReLU()
        (3): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (4): GroupNorm(32, 64, eps=1e-05, affine=True)
      )
      (relu): ReLU()
    )
    (15): Upsample(scale_factor=2.0, mode='bilinear')
    (16): Conv2d(64, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (17): GroupNorm(32, 32, eps=1e-05, affine=True)
    (18): ReLU()
    (19): ResBlock(
      (block): Sequential(
        (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (1): GroupNorm(32, 32, eps=1e-05, affine=True)
        (2): ReLU()
        (3): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (4): GroupNorm(32, 32, eps=1e-05, affine=True)
      )
      (relu): ReLU()
    )
    (20): Conv2d(32, 3, kernel_size=(1, 1), stride=(1, 1))
    (21): Sigmoid()
  )
)

perceptual_loss_fn = LPIPSLoss().to(device)

print(perceptual_loss_fn)

Setting up [LPIPS] perceptual loss: trunk [vgg], v[0.1], spatial [off]

/home/kmcalist/.local/lib/python3.10/site-packages/torchvision/models/_utils.py:208: UserWarning:

The parameter 'pretrained' is deprecated since 0.13 and may be removed in the future, please use 'weights' instead.

/home/kmcalist/.local/lib/python3.10/site-packages/torchvision/models/_utils.py:223: UserWarning:

Arguments other than a weight enum or `None` for 'weights' are deprecated since 0.13 and may be removed in the future. The current behavior is equivalent to passing `weights=VGG16_Weights.IMAGENET1K_V1`. You can also use `weights=VGG16_Weights.DEFAULT` to get the most up-to-date weights.

Loading model from: /home/kmcalist/.local/lib/python3.10/site-packages/lpips/weights/v0.1/vgg.pth
LPIPSLoss(
  (lpips): LPIPS(
    (scaling_layer): ScalingLayer()
    (net): vgg16(
      (slice1): Sequential(
        (0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (1): ReLU(inplace=True)
        (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (3): ReLU(inplace=True)
      )
      (slice2): Sequential(
        (4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
        (5): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (6): ReLU(inplace=True)
        (7): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (8): ReLU(inplace=True)
      )
      (slice3): Sequential(
        (9): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
        (10): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (11): ReLU(inplace=True)
        (12): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (13): ReLU(inplace=True)
        (14): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (15): ReLU(inplace=True)
      )
      (slice4): Sequential(
        (16): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
        (17): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (18): ReLU(inplace=True)
        (19): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (20): ReLU(inplace=True)
        (21): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (22): ReLU(inplace=True)
      )
      (slice5): Sequential(
        (23): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
        (24): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (25): ReLU(inplace=True)
        (26): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (27): ReLU(inplace=True)
        (28): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (29): ReLU(inplace=True)
      )
    )
    (lin0): NetLinLayer(
      (model): Sequential(
        (0): Dropout(p=0.5, inplace=False)
        (1): Conv2d(64, 1, kernel_size=(1, 1), stride=(1, 1), bias=False)
      )
    )
    (lin1): NetLinLayer(
      (model): Sequential(
        (0): Dropout(p=0.5, inplace=False)
        (1): Conv2d(128, 1, kernel_size=(1, 1), stride=(1, 1), bias=False)
      )
    )
    (lin2): NetLinLayer(
      (model): Sequential(
        (0): Dropout(p=0.5, inplace=False)
        (1): Conv2d(256, 1, kernel_size=(1, 1), stride=(1, 1), bias=False)
      )
    )
    (lin3): NetLinLayer(
      (model): Sequential(
        (0): Dropout(p=0.5, inplace=False)
        (1): Conv2d(512, 1, kernel_size=(1, 1), stride=(1, 1), bias=False)
      )
    )
    (lin4): NetLinLayer(
      (model): Sequential(
        (0): Dropout(p=0.5, inplace=False)
        (1): Conv2d(512, 1, kernel_size=(1, 1), stride=(1, 1), bias=False)
      )
    )
    (lins): ModuleList(
      (0): NetLinLayer(
        (model): Sequential(
          (0): Dropout(p=0.5, inplace=False)
          (1): Conv2d(64, 1, kernel_size=(1, 1), stride=(1, 1), bias=False)
        )
      )
      (1): NetLinLayer(
        (model): Sequential(
          (0): Dropout(p=0.5, inplace=False)
          (1): Conv2d(128, 1, kernel_size=(1, 1), stride=(1, 1), bias=False)
        )
      )
      (2): NetLinLayer(
        (model): Sequential(
          (0): Dropout(p=0.5, inplace=False)
          (1): Conv2d(256, 1, kernel_size=(1, 1), stride=(1, 1), bias=False)
        )
      )
      (3-4): 2 x NetLinLayer(
        (model): Sequential(
          (0): Dropout(p=0.5, inplace=False)
          (1): Conv2d(512, 1, kernel_size=(1, 1), stride=(1, 1), bias=False)
        )
      )
    )
  )
)

from torch import optim

# Optimization
# AdamW with default weight decay
optimizer = optim.AdamW(model.parameters(), lr=1e-3)
#Long burn
num_epochs = 500


# Cosine w/ Warm Restarts
# scheduler = optim.lr_scheduler.CosineAnnealingWarmRestarts(
#     optimizer,
#     T_0=50,          # First cycle is 50 epochs
#     T_mult=2,        # Double the length of each subsequent cycle
#     eta_min=1e-6     # Let it decay very close to zero
# )

Now, let's train the thing!

from tqdm import tqdm
import numpy as np
import torch
from plotly.subplots import make_subplots
import ipywidgets as widgets
from IPython.display import display
import base64, io
from PIL import Image as PILImage

import plotly.graph_objects as go

# =============================================================================
# DASHBOARD SETUP  — Plotly + ipywidgets (survives nbconvert → HTML)
# =============================================================================

history = {'epoch': [], 'total': [], 'nll': [], 'kld': [], 'perceptual': [], 'p_weight': []}
snapshot_store = {}   # epoch → list of 30 PIL images

# ── Loss figure (Plotly FigureWidget — updates in-place) ──────────
loss_fig = make_subplots(
    rows=1, cols=5,
    subplot_titles=('Total ELBO Loss', 'NLL  (Reconstruction)', 'KL Divergence',
                    'Perceptual Loss', 'Perceptual Weight'),
)
loss_fig.add_trace(go.Scatter(x=[], y=[], mode='lines',
                               line=dict(color='#e94560', width=2),
                               name='Total'), row=1, col=1)
loss_fig.add_trace(go.Scatter(x=[], y=[], mode='lines',
                               line=dict(color='#a0c4ff', width=2),
                               name='NLL'),   row=1, col=2)
loss_fig.add_trace(go.Scatter(x=[], y=[], mode='lines',
                               line=dict(color='#ffd700', width=2),
                               name='KLD'),   row=1, col=3)
loss_fig.add_trace(go.Scatter(x=[], y=[], mode='lines',
                               line=dict(color='#90ee90', width=2),
                               name='Perceptual'), row=1, col=4)
loss_fig.add_trace(go.Scatter(x=[], y=[], mode='lines',
                               line=dict(color='#ff9f40', width=2),
                               name='P-Weight'), row=1, col=5)
loss_fig.update_layout(
    height=320, paper_bgcolor='#1a1a2e', plot_bgcolor='#0d0d1a',
    font=dict(color='white'), showlegend=False,
    margin=dict(t=40, b=30, l=40, r=20),
)
for i in range(1, 6):
    loss_fig.update_xaxes(gridcolor='#222244', title_text='Epoch',
                           row=1, col=i)
    loss_fig.update_yaxes(gridcolor='#222244', row=1, col=i)

loss_widget = go.FigureWidget(loss_fig)


def _update_loss_plots():
    ep  = history['epoch']
    with loss_widget.batch_update():
        loss_widget.data[0].x = ep;  loss_widget.data[0].y = history['total']
        loss_widget.data[1].x = ep;  loss_widget.data[1].y = history['nll']
        loss_widget.data[2].x = ep;  loss_widget.data[2].y = history['kld']
        loss_widget.data[3].x = ep;  loss_widget.data[3].y = history['perceptual']
        loss_widget.data[4].x = ep;  loss_widget.data[4].y = history['p_weight']


# ── Snapshot viewer (ipywidgets — buttons + image grid) ───────────
def _tensor_to_b64(img_np):
    """Convert H×W×3 float32 numpy array → base64 PNG string."""
    arr = (np.clip(img_np, 0, 1) * 255).astype(np.uint8)
    buf = io.BytesIO()
    PILImage.fromarray(arr).save(buf, format='PNG')
    return 'data:image/png;base64,' + base64.b64encode(buf.getvalue()).decode()

ROW_LABELS = ['Train  Orig / Recon', 'Val  Orig / Recon', 'Generated']

# 30 image widgets arranged in 3 rows × 10 cols
img_widgets = [[widgets.Image(format='png', width=80, height=80)
                for _ in range(10)] for _ in range(3)]

row_boxes = []
for r in range(3):
    label = widgets.Label(
        value=ROW_LABELS[r],
        layout=widgets.Layout(width='130px', display='flex',
                              align_items='center'))
    row_box = widgets.HBox(
        [label] + img_widgets[r],
        layout=widgets.Layout(align_items='center'))
    row_boxes.append(row_box)

snap_label = widgets.Label(
    value='Snapshot — no data yet',
    layout=widgets.Layout(margin='4px 0 4px 0'))
snap_label.style.font_size = '13px'

snap_idx = [0]

btn_prev = widgets.Button(description='◀  Prev',
                          button_style='',
                          layout=widgets.Layout(width='100px'))
btn_next = widgets.Button(description='Next  ▶',
                          button_style='',
                          layout=widgets.Layout(width='100px'))
for b in (btn_prev, btn_next):
    b.style.button_color = '#16213e'
    b.style.text_color   = 'white'

epoch_dropdown = widgets.Dropdown(
    options=[],
    description='Epoch:',
    layout=widgets.Layout(width='150px'))


def _render_snapshot(epoch):
    imgs = snapshot_store.get(epoch)
    if imgs is None:
        return
    for r in range(3):
        for c in range(10):
            b64 = _tensor_to_b64(imgs[r * 10 + c])
            raw = base64.b64decode(b64.split(',')[1])
            img_widgets[r][c].value = raw
    snap_label.value = f'Snapshot — Epoch {epoch}'


def _on_prev(_b):
    keys = sorted(snapshot_store.keys())
    if not keys: return
    snap_idx[0] = max(0, snap_idx[0] - 1)
    epoch_dropdown.value = keys[snap_idx[0]]


def _on_next(_b):
    keys = sorted(snapshot_store.keys())
    if not keys: return
    snap_idx[0] = min(len(keys) - 1, snap_idx[0] + 1)
    epoch_dropdown.value = keys[snap_idx[0]]


def _on_dropdown_change(change):
    if change['name'] == 'value' and change['new'] is not None:
        keys = sorted(snapshot_store.keys())
        snap_idx[0] = keys.index(change['new'])
        _render_snapshot(change['new'])


btn_prev.on_click(_on_prev)
btn_next.on_click(_on_next)
epoch_dropdown.observe(_on_dropdown_change)

nav_bar      = widgets.HBox([btn_prev, epoch_dropdown, btn_next])
snap_viewer  = widgets.VBox([snap_label, nav_bar] + row_boxes)
dashboard    = widgets.VBox([loss_widget, snap_viewer])
display(dashboard)


def _capture_snapshot(epoch):
    model.eval()
    imgs = []
    with torch.no_grad():
        tb, _ = next(iter(train_loader))
        tb = tb[:5].to(device);  tr, _, _ = model(tb)
        for i in range(5):
            imgs.append(tb[i].cpu().permute(1,2,0).numpy())
            imgs.append(tr[i].cpu().permute(1,2,0).numpy())

        vb, _ = next(iter(val_loader))
        vb = vb[:5].to(device);  vr, _, _ = model(vb)
        for i in range(5):
            imgs.append(vb[i].cpu().permute(1,2,0).numpy())
            imgs.append(vr[i].cpu().permute(1,2,0).numpy())

        z_rand = torch.randn(10, hidden_size).to(device) * 0.7
        gen = model.decoder(model.fc_decode(z_rand).view(-1, 512, 4, 4))
        for i in range(10):
            imgs.append(gen[i].cpu().permute(1,2,0).numpy())

    snapshot_store[epoch] = imgs
    epoch_dropdown.options = sorted(snapshot_store.keys())
    snap_idx[0] = len(snapshot_store) - 1
    epoch_dropdown.value = epoch
    model.train()

# =============================================================================
# TRAINING LOOP  — VAE + Perceptual Loss
# =============================================================================

best_avg_loss = float('inf')
global_step   = 0

for epoch in range(1, num_epochs + 1):
    model.train()
    train_loss    = 0.0
    total_nll     = 0.0
    total_kld     = 0.0
    total_p_loss  = 0.0
    total_p_weight = 0.0

    #beta = 1.0
    beta = get_beta(epoch, warmup_start=5, warmup_end=15)
    pbar = tqdm(train_loader,
                desc=f"Epoch {epoch:3d}/{num_epochs}", leave=False)

    for batch_idx, (data, _) in enumerate(pbar):
        data = data.to(device)

        optimizer.zero_grad()
        recon_batch, mu, logvar = model(data)

        last_layer_weight = model.get_last_decoder_layer()

        loss, nll, kld, p_loss, p_weight = perceptual_generator_loss(
            reconstruction     = recon_batch,
            x                  = data,
            mu                 = mu,
            logvar             = logvar,
            perceptual_loss_fn = perceptual_loss_fn,
            last_layer_weight  = last_layer_weight,
            beta               = beta,
            free_bits          = 0.1,
        )

        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=10.0)
        optimizer.step()

        train_loss     += loss.item()
        total_nll      += nll
        total_kld      += kld
        total_p_loss   += p_loss
        total_p_weight += p_weight
        global_step    += 1

        pbar.set_postfix({
            'Loss':   f'{train_loss    / (batch_idx+1):.1f}',
            'NLL':    f'{total_nll     / (batch_idx+1):.1f}',
            'KLD':    f'{total_kld     / (batch_idx+1):.1f}',
            'PL':     f'{total_p_loss  / (batch_idx+1):.4f}',
            'PW':     f'{total_p_weight/ (batch_idx+1):.2f}',
            'β':      f'{beta:.3f}',
        })

    # ── Epoch averages ────────────────────────────────────────────
    n              = len(train_loader)
    avg_loss       = train_loss    / n
    avg_nll        = total_nll     / n
    avg_kld        = total_kld     / n
    avg_p_loss     = total_p_loss  / n
    avg_p_weight   = total_p_weight / n
    current_lr     = optimizer.param_groups[0]['lr']

    # ── Checkpoint on best training ELBO ─────────────────────────
    if avg_loss < best_avg_loss:
        best_avg_loss = avg_loss
        torch.save({
            "epoch":            epoch,
            "global_step":      global_step,
            "model_state_dict": model.state_dict(),
            "opt_state_dict":   optimizer.state_dict(),
            "best_avg_loss":    best_avg_loss,
            "hidden_size":      model.hidden_size,
        }, "vae_perceptual_best.pth")

    #scheduler.step()

    # ── Live dashboard update ─────────────────────────────────────
    if epoch >= 20:
        history['epoch']     .append(epoch)
        history['total']     .append(avg_loss)
        history['nll']       .append(avg_nll)
        history['kld']       .append(avg_kld)
        history['perceptual'].append(avg_p_loss)
        history['p_weight']  .append(avg_p_weight)
        _update_loss_plots()

    # ── Snapshot every 10 epochs ──────────────────────────────────
    if epoch % 10 == 0:
        _capture_snapshot(epoch)

VBox(children=(FigureWidget({
    'data': [{'line': {'color': '#e94560', 'width': 2},
              'mode': 'l…

An Improvement: The Two Stage Sampler

Above we can see that our reconstructions are decent now! Not blurry. Just... off a little. Facial features are weird, edges are a little weird. There's still a little blur.

That all said, the reconstructions are recognizable, but the generated images look like dog and cat parts thrown together!

Our theory is that each $\mathbf X$ is a draw of the random vector $\mathbf z$ from $P(Z)$. $P(Z)$ dictates the arrangment of the 256 latent values that are coherent - the VAE is learning a mapping/projection of the data space into the cat/dog space and vice-versa.

The problem is that we don't know $P(Z)$! We've learned $P(z | \mathbf X)$, but not the marginal distribution. With a large enough training set, we expect that $P(\mathbf z | \mathbf X)$ represents the true latent distribution - correlations and all - since we've seen a lot of $\mathbf X$ values and mapped them to optimal $\mathbf z$ values.

But, do we know the structure of $P(\mathbf z | \mathbf X)$? Can we sample from it?

By construction:

$$ P(\mathbf z | \mathbf X) = P(\mathbf X | f_\theta(\mathbf z)) \mathcal N_K(\mathbf z | \mathbf 0 , \mathbf I) $$

So, we expect that the posterior is close-ish to a normal distribution centered on 0 and with unit variance.

What if we take draws from this distribution and pass them through the decoder?

import matplotlib.pyplot as plt

model.eval()
with torch.no_grad():
    # Draw 100 random samples from the normal prior N(0, I)
    z_random = torch.randn(100, hidden_size).to(device)
    
    # Project and decode
    z_projected = model.fc_decode(z_random).view(-1, 512, 4, 4)
    generated_images = model.decoder(z_projected)

# Plot a 10x10 grid of the generated images
plt.figure(figsize=(15, 15))
for i in range(100):
    # Convert from (C, H, W) to (H, W, C) for plotting
    img_np = generated_images[i].cpu().permute(1, 2, 0).numpy()
    
    plt.subplot(10, 10, i + 1)
    plt.imshow(img_np)
    plt.axis('off')

plt.tight_layout()
plt.show()

It's close! But prior $\neq$ posterior.

Better approach - try to learn the density of $\mathbf z | x$. At its core, this is what a VAE does - learn a low-dimensional representation (a proper PDF) of latent values that describe $\mathbf X$.

So why not train a second, smaller VAE on the encoded latent vectors themselves?

Two-Stage Sampling

After training the first VAE, we encode the entire training set and collect the posterior means $\{\mu_1, \mu_2, \ldots, \mu_N\}$. These are the points in latent space where real images actually live. The KL penalty pulled them toward $\mathcal{N}(0, I)$, but they don't match it exactly — there are regions of high prior density where no real images map, and clusters of real images that the prior underweights.

The second VAE learns exactly this distribution. It takes the $\mu$ vectors as input, compresses them into an even smaller latent space $\mathbf{z}_2$, and learns to reconstruct them. Its own KL penalty ensures that its latent space is well-structured and sampleable.

The generation pipeline then becomes:

1. Sample $\mathbf{z}_2 \sim \mathcal{N}(0, I)$ from the second VAE's prior
2. Decode through the second VAE to get $\hat{\mu}$ — a point in the first VAE's latent space that lies on the manifold where real images actually encode
3. Decode $\hat{\mu}$ through the original VAE's decoder to get an image

The second VAE's job is much easier than the first. The distribution of $\mu$ vectors is already roughly Gaussian (the KL penalty ensured that), low-dimensional, and smooth. A small network with a few linear layers trains in minutes and closes the prior-posterior gap almost entirely.

Key detail: the two stages must be trained sequentially, not jointly. The first VAE establishes good representations and a good decoder. The second VAE corrects the sampling distribution without touching either. Joint training degrades both — the first VAE's latent space shifts to accommodate the second VAE's preferences, which destabilizes the decoder's learned mapping (Dai & Wipf, 2019).

The result: generations that cover the full diversity of the training data instead of clustering around the highest-density regions of the prior. Different breeds, poses, colorings, and backgrounds — all because we're now sampling from where the data actually lives rather than where we assumed it would be.

# =============================================================================
# STAGE 2 VAE — learns the distribution of latent vectors from Stage 1
# =============================================================================

class VAE2(nn.Module):
    """Small MLP VAE that models the distribution of Stage-1 mu vectors."""
    def __init__(self, input_dim=256, hidden_dim=512, latent_dim=64):
        super().__init__()
        self.latent_dim = latent_dim

        self.encoder = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
        )
        self.fc_mu     = nn.Linear(hidden_dim, latent_dim)
        self.fc_logvar = nn.Linear(hidden_dim, latent_dim)

        self.decoder = nn.Sequential(
            nn.Linear(latent_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, input_dim),
        )

    def reparameterize(self, mu, logvar):
        std = torch.exp(0.5 * logvar)
        return mu + std * torch.randn_like(std)

    def forward(self, z):
        h      = self.encoder(z)
        mu     = self.fc_mu(h)
        logvar = torch.clamp(self.fc_logvar(h), -10, 10)
        z2     = self.reparameterize(mu, logvar)
        recon  = self.decoder(z2)
        return recon, mu, logvar

# ── 1. Collect Stage-1 posterior means over the full training set ──
model.eval()
all_mus = []
with torch.no_grad():
    for data, _ in train_loader:
        data   = data.to(device)
        enc    = model.encoder(data)
        mu_vec = model.fc_mu(enc)
        all_mus.append(mu_vec.cpu())

all_mus = torch.cat(all_mus, dim=0)   # (N, 256)
print(f"Collected {all_mus.shape[0]} latent vectors, dim={all_mus.shape[1]}")

Collected 12500 latent vectors, dim=256

# ── 2. Build a simple Dataset / DataLoader for Stage-2 training ──
latent_dataset = torch.utils.data.TensorDataset(all_mus)
latent_loader  = torch.utils.data.DataLoader(
    latent_dataset, batch_size=512, shuffle=True, num_workers=4)

# ── 3. Instantiate Stage-2 VAE and optimizer ──────────────────────
vae2           = VAE2(input_dim=hidden_size, hidden_dim=512, latent_dim=128).to(device)
optimizer2     = optim.AdamW(vae2.parameters(), lr=1e-3)
num_epochs2    = 1000
best2_loss     = float('inf')

print(vae2)

# ── 4. Training loop ──────────────────────────────────────────────
history2 = {'epoch': [], 'total': [], 'nll': [], 'kld': []}

loss_fig2 = make_subplots(rows=1, cols=3,
    subplot_titles=('Total Loss', 'NLL', 'KLD'))
for col, name, color in zip([1,2,3],
                              ['Total','NLL','KLD'],
                              ['#e94560','#a0c4ff','#ffd700']):
    loss_fig2.add_trace(
        go.Scatter(x=[], y=[], mode='lines',
                   line=dict(color=color, width=2), name=name),
        row=1, col=col)
loss_fig2.update_layout(
    height=280, paper_bgcolor='#1a1a2e', plot_bgcolor='#0d0d1a',
    font=dict(color='white'), showlegend=False,
    margin=dict(t=40, b=30, l=40, r=20))
loss_widget2 = go.FigureWidget(loss_fig2)
display(loss_widget2)

for epoch in range(1, num_epochs2 + 1):
    vae2.train()
    ep_total = ep_nll = ep_kld = 0.0

    for (z_batch,) in latent_loader:
        z_batch = z_batch.to(device)
        optimizer2.zero_grad()

        recon, mu2, logvar2 = vae2(z_batch)

        nll  = F.mse_loss(recon, z_batch, reduction='sum') / z_batch.shape[0]
        kld  = (-0.5 * (1 + logvar2 - mu2.pow(2) - logvar2.exp())
                ).sum() / z_batch.shape[0]
        loss = nll + (0.5*kld)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(vae2.parameters(), 5.0)
        optimizer2.step()

        ep_total += loss.item()
        ep_nll   += nll.item()
        ep_kld   += kld.item()

    n = len(latent_loader)
    if epoch % 5 == 0:
        history2['epoch'].append(epoch)
        history2['total'].append(ep_total / n)
        history2['nll']  .append(ep_nll   / n)
        history2['kld']  .append(ep_kld   / n)
        with loss_widget2.batch_update():
            for i, key in enumerate(['total','nll','kld']):
                loss_widget2.data[i].x = history2['epoch']
                loss_widget2.data[i].y = history2[key]

    if ep_total / n < best2_loss:
        best2_loss = ep_total / n
        torch.save(vae2.state_dict(), "vae2_best.pth")

    if epoch % 100 == 0:
        print(f"Epoch {epoch:3d}/{num_epochs2}  "
              f"total={ep_total/n:.2f}  nll={ep_nll/n:.2f}  kld={ep_kld/n:.2f}")

VAE2(
  (encoder): Sequential(
    (0): Linear(in_features=256, out_features=512, bias=True)
    (1): ReLU()
    (2): Linear(in_features=512, out_features=512, bias=True)
    (3): ReLU()
  )
  (fc_mu): Linear(in_features=512, out_features=128, bias=True)
  (fc_logvar): Linear(in_features=512, out_features=128, bias=True)
  (decoder): Sequential(
    (0): Linear(in_features=128, out_features=512, bias=True)
    (1): ReLU()
    (2): Linear(in_features=512, out_features=512, bias=True)
    (3): ReLU()
    (4): Linear(in_features=512, out_features=256, bias=True)
  )
)

FigureWidget({
    'data': [{'line': {'color': '#e94560', 'width': 2},
              'mode': 'lines',
              'name': 'Total',
              'type': 'scatter',
              'uid': 'acad75eb-28eb-4c70-9d1c-4c8bd363c595',
              'x': [],
              'xaxis': 'x',
              'y': [],
              'yaxis': 'y'},
             {'line': {'color': '#a0c4ff', 'width': 2},
              'mode': 'lines',
              'name': 'NLL',
              'type': 'scatter',
              'uid': '7a54c580-db21-4e3a-bdb2-18cd006db9f7',
              'x': [],
              'xaxis': 'x2',
              'y': [],
              'yaxis': 'y2'},
             {'line': {'color': '#ffd700', 'width': 2},
              'mode': 'lines',
              'name': 'KLD',
              'type': 'scatter',
              'uid': '2f9816e1-bd8b-4222-bf24-53f0a2813593',
              'x': [],
              'xaxis': 'x3',
              'y': [],
              'yaxis': 'y3'}],
    'layout': {'annotations': [{'font': {'size': 16},
                                'showarrow': False,
                                'text': 'Total Loss',
                                'x': 0.14444444444444446,
                                'xanchor': 'center',
                                'xref': 'paper',
                                'y': 1.0,
                                'yanchor': 'bottom',
                                'yref': 'paper'},
                               {'font': {'size': 16},
                                'showarrow': False,
                                'text': 'NLL',
                                'x': 0.5,
                                'xanchor': 'center',
                                'xref': 'paper',
                                'y': 1.0,
                                'yanchor': 'bottom',
                                'yref': 'paper'},
                               {'font': {'size': 16},
                                'showarrow': False,
                                'text': 'KLD',
                                'x': 0.8555555555555556,
                                'xanchor': 'center',
                                'xref': 'paper',
                                'y': 1.0,
                                'yanchor': 'bottom',
                                'yref': 'paper'}],
               'font': {'color': 'white'},
               'height': 280,
               'margin': {'b': 30, 'l': 40, 'r': 20, 't': 40},
               'paper_bgcolor': '#1a1a2e',
               'plot_bgcolor': '#0d0d1a',
               'showlegend': False,
               'template': '...',
               'xaxis': {'anchor': 'y', 'domain': [0.0, 0.2888888888888889]},
               'xaxis2': {'anchor': 'y2', 'domain': [0.35555555555555557, 0.6444444444444445]},
               'xaxis3': {'anchor': 'y3', 'domain': [0.7111111111111111, 1.0]},
               'yaxis': {'anchor': 'x', 'domain': [0.0, 1.0]},
               'yaxis2': {'anchor': 'x2', 'domain': [0.0, 1.0]},
               'yaxis3': {'anchor': 'x3', 'domain': [0.0, 1.0]}}
})

Epoch 100/1000  total=145.42  nll=128.47  kld=33.90
Epoch 200/1000  total=137.72  nll=115.74  kld=43.96
Epoch 300/1000  total=134.28  nll=109.73  kld=49.09
Epoch 400/1000  total=132.56  nll=106.75  kld=51.62
Epoch 500/1000  total=131.56  nll=104.99  kld=53.15
Epoch 600/1000  total=130.76  nll=103.75  kld=54.02
Epoch 700/1000  total=130.24  nll=102.79  kld=54.91
Epoch 800/1000  total=129.84  nll=102.10  kld=55.48
Epoch 900/1000  total=129.39  nll=101.29  kld=56.21
Epoch 1000/1000  total=129.15  nll=100.74  kld=56.82

# ── 5. Two-stage generation: N(0,I) → VAE2 decoder → VAE1 decoder ─
vae2.eval()
model.eval()

with torch.no_grad():
    z2_sample   = torch.randn(100, vae2.latent_dim).to(device)
    z1_hat      = vae2.decoder(z2_sample)                          # (100, 256)
    z_proj      = model.fc_decode(z1_hat).view(-1, 512, 4, 4)
    gen_images  = model.decoder(z_proj)                            # (100, 3, 64, 64)

plt.figure(figsize=(15, 15))
for i in range(100):
    img_np = gen_images[i].cpu().permute(1, 2, 0).numpy()
    plt.subplot(10, 10, i + 1)
    plt.imshow(img_np)
    plt.axis('off')
plt.suptitle('Two-Stage VAE Generations: N(0,I) → VAE2 → VAE1', color='white')
plt.tight_layout()
plt.show()

A little better...

That's not just blurry cat outlines - it's definitely cats and dogs. Just weird looking...

We'll fix that soon.

Why VAEs

Two big reasons that we like VAEs

1. Images are adjustible using latent vector math

2. Easyish to conditionally generate

Let's start by demonstrating how we can use embeddings to train a conditioned VAE.

Suppose that we have features that are associated with each image. In this case, we know whether an image is a cat or a dog. Suppose that I wanted to use my VAE to generate a cat. I could create two separate VAEs, but this would cause me to miss out on the shared information that can be learned about animal structure from both cats and dogs.

Instead, we'll create a categorical embedding of length D for each class. As we've seen previously, this is a fully differentiable way to add categorical features to a neural network.

Conditional VAEs

We have cats and dogs. Currently, we are just encoding the images into the latent vector space (e.g. the continuous latent codebook) and letting the decoder learn how similar images should relate to one another in this space. But, this isn't all the information that we have! First, we know whether an image is a cat or a dog. Additionally, we can use pre-trained classifiers trained on ImageNet to learn something about the likely breed of the dog or cat. Finally, we can use instance segmentation to find the pixels that correspond to cat or dog and figure out the main color of the critter in the image. These are important pieces of information that we should give our VAE.

As stated above, though, we don't want to do this separately because there is shared information between categories - a brown dog and a brown cat are both brown while a golden retriever is still a dog even if its features aren't identical to a corgi.

# """
# Animal Labeling Pipeline — Notebook Version (Fast)
# ===================================================
# For each image in dogs-vs-cats/train, extract:
#   - label (cat/dog from filename)
#   - breed (EfficientNet-B0, filtered to species-appropriate ImageNet classes)
#   - color (YOLO11 instance segmentation + vectorized pixel classification)

# Outputs: animal_labels.csv with columns [filename, label, breed, color]

# Requirements:
#   pip install ultralytics scikit-learn pandas tqdm
# """

# import os
# import glob
# import numpy as np
# import pandas as pd
# from PIL import Image
# from tqdm.notebook import tqdm

# import torch
# from torchvision import transforms
# from torchvision.models import efficientnet_b0, EfficientNet_B0_Weights
# from torch.utils.data import Dataset, DataLoader
# from ultralytics import YOLO

# # =============================================================================
# # CONFIG — change these as needed
# # =============================================================================
# DATA_DIR = "dogs-vs-cats/train"
# OUTPUT_CSV = "animal_labels.csv"
# LIMIT = None            # Set to e.g. 100 for testing, None for full dataset
# BREED_BATCH_SIZE = 64   # Batch size for EfficientNet
# SEG_BATCH_SIZE = 32     # Batch size for YOLO segmentation
# NUM_WORKERS = 4         # DataLoader workers

# # =============================================================================
# # DEVICE
# # =============================================================================
# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# print(f"Using device: {device}")

# # =============================================================================
# # IMAGENET CAT AND DOG BREED INDICES
# # =============================================================================
# DOG_BREED_INDICES = set(range(151, 269))

# CAT_BREED_INDICES = {
#     281,  # tabby
#     282,  # tiger cat
#     283,  # Persian cat
#     284,  # Siamese cat
#     285,  # Egyptian cat
# }

# # =============================================================================
# # LOAD MODELS
# # =============================================================================
# print("Loading EfficientNet-B0 for breed classification...")
# breed_weights = EfficientNet_B0_Weights.IMAGENET1K_V1
# breed_model = efficientnet_b0(weights=breed_weights).eval().to(device)
# imagenet_classes = breed_weights.meta["categories"]

# breed_preprocess = transforms.Compose([
#     transforms.Resize(256),
#     transforms.CenterCrop(224),
#     transforms.ToTensor(),
#     transforms.Normalize(mean=[0.485, 0.456, 0.406],
#                          std=[0.229, 0.224, 0.225]),
# ])

# dog_mask = torch.zeros(len(imagenet_classes), dtype=torch.bool, device=device)
# cat_mask = torch.zeros(len(imagenet_classes), dtype=torch.bool, device=device)
# for i in DOG_BREED_INDICES:
#     dog_mask[i] = True
# for i in CAT_BREED_INDICES:
#     cat_mask[i] = True

# print(f"  Dog breed classes: {dog_mask.sum().item()}")
# print(f"  Cat breed classes: {cat_mask.sum().item()}")

# print("Loading YOLO11 for instance segmentation...")
# yolo = YOLO("yolo11m-seg.pt")

# YOLO_CAT = 15
# YOLO_DOG = 16

# print("Models loaded.\n")


# # =============================================================================
# # STAGE 1: BATCHED BREED CLASSIFICATION
# # =============================================================================

# BREED_CONF_THRESHOLD = 0.5   # ← tune this; lower = more specific breeds assigned

# class BreedDataset(Dataset):
#     def __init__(self, image_paths, transform):
#         self.image_paths = image_paths
#         self.transform = transform

#     def __len__(self):
#         return len(self.image_paths)

#     def __getitem__(self, idx):
#         try:
#             img = Image.open(self.image_paths[idx]).convert("RGB")
#             return self.transform(img), idx
#         except Exception:
#             return torch.zeros(3, 224, 224), idx


# def classify_breeds_batched(image_paths, labels):
#     """Classify breeds in batches, filtered by species."""
#     dataset = BreedDataset(image_paths, breed_preprocess)
#     loader = DataLoader(
#         dataset,
#         batch_size=BREED_BATCH_SIZE,
#         shuffle=False,
#         num_workers=NUM_WORKERS,
#         pin_memory=True,
#     )

#     breeds = ["unknown"] * len(image_paths)

#     with torch.no_grad():
#         for batch_imgs, batch_indices in tqdm(loader, desc="Stage 1: Breed classification"):
#             batch_imgs = batch_imgs.to(device)
#             logits = breed_model(batch_imgs)
#             probs = torch.softmax(logits, dim=1)

#             for i, idx in enumerate(batch_indices):
#                 idx = idx.item()
#                 is_dog = labels[idx] == "dog"
#                 p = probs[i].clone()

#                 # Zero out wrong-species classes
#                 p[~dog_mask if is_dog else ~cat_mask] = 0.0

#                 # Re-normalise so confidences sum to 1 within species
#                 p_sum = p.sum()
#                 if p_sum < 1e-8:
#                     # EfficientNet had no species-matching activation at all
#                     breeds[idx] = "OtherDog" if is_dog else "OtherCat"
#                     continue

#                 p = p / p_sum   # normalised within-species probability

#                 top_prob = p.max().item()
#                 if top_prob >= BREED_CONF_THRESHOLD:
#                     breeds[idx] = imagenet_classes[p.argmax().item()]
#                 else:
#                     # Model is uncertain — assign generic fallback
#                     breeds[idx] = "OtherDog" if is_dog else "OtherCat"

#     return breeds


# # =============================================================================
# # STAGE 2: YOLO SEGMENTATION + VECTORIZED COLOR CLASSIFICATION
# # =============================================================================

# COLOR_NAMES = ["black", "gray", "white", "cream", "golden", "ginger", "brown"]

# def get_dominant_color_vectorized(img_np, mask):
#     # Resize mask to image dimensions if needed
#     if mask.shape != img_np.shape[:2]:
#         mask = np.array(
#             Image.fromarray((mask * 255).astype(np.uint8)).resize(
#                 (img_np.shape[1], img_np.shape[0]), Image.BILINEAR
#             )
#         ) > 64

#     animal_pixels = img_np[mask].astype(float)
#     if len(animal_pixels) < 50:
#         return "unknown"

#     # ── No subsampling — vectorized ops are fast on all pixels ───
#     r_n = animal_pixels[:, 0] / 255.0
#     g_n = animal_pixels[:, 1] / 255.0
#     b_n = animal_pixels[:, 2] / 255.0

#     cmax = np.maximum(np.maximum(r_n, g_n), b_n)
#     cmin = np.minimum(np.minimum(r_n, g_n), b_n)
#     diff = cmax - cmin

#     h = np.zeros_like(cmax)
#     nonzero = diff > 0
#     mask_r = (cmax == r_n) & nonzero
#     mask_g = (cmax == g_n) & nonzero
#     mask_b = (cmax == b_n) & nonzero
#     h[mask_r] = (60 * ((g_n[mask_r] - b_n[mask_r]) / diff[mask_r]) + 360) % 360
#     h[mask_g] = (60 * ((b_n[mask_g] - r_n[mask_g]) / diff[mask_g]) + 120) % 360
#     h[mask_b] = (60 * ((r_n[mask_b] - g_n[mask_b]) / diff[mask_b]) + 240) % 360

#     s = np.where(cmax > 0, diff / cmax, 0)
#     v = cmax

#     # ── Pet-appropriate color labels ──────────────────────────────
#     # 0=black  1=gray  2=white  3=cream  4=golden  5=ginger  6=brown
#     pixel_labels = np.full(len(animal_pixels), -1, dtype=int)

#     # Very dark → black regardless of hue
#     pixel_labels[v < 0.20] = 0

#     # Low saturation (achromatic)
#     achromatic = (s < 0.18) & (pixel_labels == -1)
#     pixel_labels[achromatic & (v < 0.45)]               = 1  # gray
#     pixel_labels[achromatic & (v >= 0.45) & (v < 0.80)] = 3  # cream
#     pixel_labels[achromatic & (v >= 0.80)]               = 2  # white

#     # Chromatic
#     chroma = (s >= 0.18) & (v >= 0.20) & (pixel_labels == -1)
#     pixel_labels[chroma & ((h < 20) | (h >= 345))]      = 5  # ginger/red
#     pixel_labels[chroma & (h >= 20) & (h < 75)]         = 4  # golden/orange-yellow
#     pixel_labels[chroma & (h >= 75) & (pixel_labels == -1)] = 6  # brown

#     pixel_labels[pixel_labels == -1] = 6  # catch-all

#     counts = np.bincount(pixel_labels, minlength=len(COLOR_NAMES))
#     return COLOR_NAMES[counts.argmax()]


# def segment_and_color_batched(image_paths):
#     colors = ["unknown"] * len(image_paths)

#     for batch_start in tqdm(range(0, len(image_paths), SEG_BATCH_SIZE),
#                             desc="Stage 2: Segmentation + color"):
#         batch_end   = min(batch_start + SEG_BATCH_SIZE, len(image_paths))
#         batch_paths = image_paths[batch_start:batch_end]
#         results     = yolo(batch_paths, verbose=False, classes=[YOLO_CAT, YOLO_DOG])

#         for i, result in enumerate(results):
#             global_idx = batch_start + i
#             try:
#                 img_np = np.array(Image.open(batch_paths[i]).convert("RGB"))

#                 if result.masks is not None and len(result.masks) > 0:
#                     # Union of ALL detected animal masks
#                     combined_mask = np.zeros(img_np.shape[:2], dtype=bool)
#                     for m_idx in range(len(result.masks)):
#                         seg_mask = result.masks.data[m_idx].cpu().numpy().astype(bool)
#                         if seg_mask.shape != img_np.shape[:2]:
#                             seg_mask = np.array(
#                                 Image.fromarray((seg_mask * 255).astype(np.uint8)).resize(
#                                     (img_np.shape[1], img_np.shape[0]), Image.BILINEAR
#                                 )
#                             ) > 64
#                         combined_mask |= seg_mask
#                     colors[global_idx] = get_dominant_color_vectorized(img_np, combined_mask)
#                 else:
#                     # Fallback: wide center crop
#                     h, w = img_np.shape[:2]
#                     m = 0.10
#                     center_mask = np.zeros((h, w), dtype=bool)
#                     center_mask[int(h*m):int(h*(1-m)), int(w*m):int(w*(1-m))] = True
#                     colors[global_idx] = get_dominant_color_vectorized(img_np, center_mask)

#             except Exception as e:
#                 tqdm.write(f"  Failed: {os.path.basename(batch_paths[i])} — {e}")

#     return colors


# # =============================================================================
# # RUN PIPELINE
# # =============================================================================

# # Collect images
# image_paths = sorted(glob.glob(os.path.join(DATA_DIR, "*.jpg")))
# if LIMIT:
#     image_paths = image_paths[:LIMIT]

# filenames = [os.path.basename(p) for p in image_paths]
# labels = ["dog" if fn.startswith("dog") else "cat" for fn in filenames]

# print(f"Processing {len(image_paths)} images from {DATA_DIR}")
# print(f"  Dogs: {labels.count('dog')}, Cats: {labels.count('cat')}\n")

# # Stage 1: Breed classification (batched EfficientNet)
# breeds = classify_breeds_batched(image_paths, labels)

# # Stage 2: Segmentation + color (batched YOLO + vectorized color)
# colors = segment_and_color_batched(image_paths)

# # =============================================================================
# # BUILD CSV
# # =============================================================================
# df = pd.DataFrame({
#     "filename": filenames,
#     "label": labels,
#     "breed": breeds,
#     "color": colors,
# })

# df.to_csv(OUTPUT_CSV, index=False)

# # =============================================================================
# # SUMMARY
# # =============================================================================
# print(f"\nSaved {len(df)} rows to {OUTPUT_CSV}")
# print("\n" + "=" * 50)
# print("SUMMARY")
# print("=" * 50)

# print(f"\nLabels:")
# for lbl, count in df["label"].value_counts().items():
#     print(f"  {lbl}: {count}")

# print(f"\nColors:")
# for color, count in df["color"].value_counts().items():
#     print(f"  {color}: {count}")

# print(f"\nTop 15 breeds:")
# for breed, count in df["breed"].value_counts().head(15).items():
#     print(f"  {breed}: {count}")

# print(f"\nSample rows:")
# print(df.sample(10, random_state=42).to_string(index=False))

# import pandas as pd

# df = pd.read_csv('animal_labels.csv')

# breed_counts = df['breed'].value_counts()
# common_breeds = breed_counts[breed_counts > 100]

# print(f"Number of breed categories with more than 100 instances: {len(common_breeds)}\n")
# print(common_breeds.to_string())

# color_counts = df['color'].value_counts()
# print(f"\nColor distribution:\n{color_counts.to_string()}")

# # Recode rare breeds to OtherDog or OtherCat
# threshold = 100
# rare_breeds = breed_counts[breed_counts <= threshold].index

# def recode_breed(row):
#     if row['breed'] in rare_breeds:
#         return 'OtherDog' if row['label'] == 'dog' else 'OtherCat'
#     return row['breed']

# df['breed'] = df.apply(recode_breed, axis=1)

# print(f"\nBreed distribution after recoding (threshold={threshold}):")
# print(df['breed'].value_counts().to_string())

# df.to_csv('petadata.csv', index=False)
# print(f"\nSaved {len(df)} rows to petadata.csv")

import os
import pandas as pd
os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import numpy as np
import lpips
from torch.utils.data import DataLoader, Dataset
from torchvision import transforms
from PIL import Image
import glob
import matplotlib.pyplot as plt

class DogCatDataset(Dataset):
    def __init__(self, folder_path, meta_path='petadata.csv', transform=None):
        self.image_paths = glob.glob(os.path.join(folder_path, '*.jpg'))
        self.transform   = transform

        # ── Load metadata and index by filename ──────────────────
        meta = pd.read_csv(meta_path)
        meta = meta.set_index('filename')

        # ── Build string→int codes for breed and color ───────────
        breeds = sorted(meta['breed'].unique())
        colors = sorted(meta['color'].unique())
        self.breed_to_idx = {b: i for i, b in enumerate(breeds)}
        self.color_to_idx = {c: i for i, c in enumerate(colors)}

        # Human-readable label names
        self.label_names = {0: 'cat', 1: 'dog'}
        self.breed_names  = {i: b for b, i in self.breed_to_idx.items()}
        self.color_names  = {i: c for c, i in self.color_to_idx.items()}

        # ── Per-image attributes ──────────────────────────────────
        self.labels  = []
        self.breeds  = []
        self.colors  = []

        for path in self.image_paths:
            filename = os.path.basename(path)

            # Species from filename prefix
            self.labels.append(1 if filename.startswith('dog') else 0)

            # Breed & color from CSV
            row = meta.loc[filename]
            self.breeds.append(self.breed_to_idx[row['breed']])
            self.colors.append(self.color_to_idx[row['color']])

        print(f"Loaded {len(self.image_paths)} images | "
              f"{len(breeds)} breeds | {len(colors)} colors")

    def __len__(self):
        return len(self.image_paths)

    def __getitem__(self, idx):
        image = Image.open(self.image_paths[idx]).convert('RGB')
        if self.transform:
            image = self.transform(image)
        return (
            image,
            self.labels[idx],           # species  (int)
            self.breeds[idx],           # breed    (int)
            self.colors[idx],           # color    (int)
        )

# Define standard preprocessing
train_transform = transforms.Compose([
    transforms.Resize((72, 72)),           # slightly larger than target
    transforms.RandomCrop((64, 64)),       # random 64×64 crop
    transforms.RandomHorizontalFlip(p=0.5),
    transforms.ToTensor(),
    transforms.Lambda(lambda x: x.clamp(0, 1))
])

# Validation gets no augmentation — deterministic center crop only
val_transform = transforms.Compose([
    transforms.Resize((64, 64)),
    transforms.ToTensor(),
    transforms.Lambda(lambda x: x.clamp(0, 1))
])

dog_cat_dataset_train = DogCatDataset('dogs-vs-cats/train', transform=train_transform)
dog_cat_dataset_val   = DogCatDataset('dogs-vs-cats/train', transform=val_transform)

train_size = int(0.5 * len(dog_cat_dataset_train))
val_size   = len(dog_cat_dataset_train) - train_size

# Fix the split indices so both datasets use the same images for train/val
indices = torch.randperm(len(dog_cat_dataset_train)).tolist()
train_dataset = torch.utils.data.Subset(dog_cat_dataset_train, indices[:train_size])
val_dataset   = torch.utils.data.Subset(dog_cat_dataset_val,   indices[train_size:])

Loaded 25000 images | 23 breeds | 7 colors
Loaded 25000 images | 23 breeds | 7 colors

# Get the base dataset
base = dog_cat_dataset_train

def get_attr_set(dataset_subset, attr):
    """Return the set of unique attribute values present in a subset."""
    base_dataset = dataset_subset.dataset
    attr_list = getattr(base_dataset, attr)
    return set(attr_list[global_idx] for global_idx in dataset_subset.indices)

train_species = get_attr_set(train_dataset, 'labels')
val_species   = get_attr_set(val_dataset,   'labels')

train_colors  = get_attr_set(train_dataset, 'colors')
val_colors    = get_attr_set(val_dataset,   'colors')

train_breeds  = get_attr_set(train_dataset, 'breeds')
val_breeds    = get_attr_set(val_dataset,   'breeds')

# Check coverage
missing_species = val_species - train_species
missing_colors  = val_colors  - train_colors
missing_breeds  = val_breeds  - train_breeds

idx_to_color = {v: k for k, v in base.color_to_idx.items()}
idx_to_breed = {v: k for k, v in base.breed_to_idx.items()}

print("=== Coverage Check: Val → Train ===\n")

print(f"Species  — val: {len(val_species)}, train: {len(train_species)}, "
      f"missing from train: {len(missing_species)}")
if missing_species:
    print(f"  Missing species indices: {missing_species}")

print(f"\nColors   — val: {len(val_colors)}, train: {len(train_colors)}, "
      f"missing from train: {len(missing_colors)}")
if missing_colors:
    print(f"  Missing colors: {[idx_to_color[i] for i in missing_colors]}")
else:
    print("  All val colors present in train ✓")

print(f"\nBreeds   — val: {len(val_breeds)}, train: {len(train_breeds)}, "
      f"missing from train: {len(missing_breeds)}")
if missing_breeds:
    print(f"  Missing breeds: {[idx_to_breed[i] for i in missing_breeds]}")
else:
    print("  All val breeds present in train ✓")

=== Coverage Check: Val → Train ===

Species  — val: 2, train: 2, missing from train: 0

Colors   — val: 7, train: 7, missing from train: 0
  All val colors present in train ✓

Breeds   — val: 23, train: 23, missing from train: 0
  All val breeds present in train ✓

train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True,  num_workers=6)
val_loader   = DataLoader(val_dataset,   batch_size=64, shuffle=False, num_workers=6)

print(f"Number of training images:   {len(train_dataset)}")
print(f"Number of validation images: {len(val_dataset)}")

Number of training images:   12500
Number of validation images: 12500

# ── Helper: get indices by attribute ──────────────────────────────
def get_indices_by_attr(dataset_subset, attr, value):
    """Return up to 4 indices from subset where attribute matches value."""
    results = []
    base_dataset = dataset_subset.dataset
    for subset_idx, global_idx in enumerate(dataset_subset.indices):
        if getattr(base_dataset, attr)[global_idx] == value:
            results.append(subset_idx)
        if len(results) >= 4:
            break
    return results

def show_examples(title, attr, value_map, dataset_subset, n_per_class=4):
    """Show n_per_class images for each category in value_map."""
    # value_map can be:
    #   {int_idx: str_name}  (label_names)  — iterate by idx
    #   {str_name: int_idx}  (breed/color_to_idx) — iterate by idx
    if isinstance(value_map, dict):
        # Normalise to {int_idx: display_name} regardless of direction
        first_key = next(iter(value_map))
        if isinstance(first_key, int):
            # {0: 'cat', 1: 'dog'}  — already idx → name
            idx_to_name = value_map
        else:
            # {'black': 0, ...}  — flip to idx → name
            idx_to_name = {v: k for k, v in value_map.items()}
    else:
        idx_to_name = {i: str(v) for i, v in enumerate(value_map)}

    categories = sorted(idx_to_name.keys())
    n_cats     = len(categories)

    fig, axes = plt.subplots(n_cats, n_per_class,
                             figsize=(n_per_class * 3, n_cats * 3))
    if n_cats == 1:
        axes = axes[np.newaxis, :]

    for row, idx in enumerate(categories):
        display_name = idx_to_name[idx]
        # Pass the integer idx — matches what's stored in base_dataset.breeds/colors/labels
        idxs = get_indices_by_attr(dataset_subset, attr, idx)
        for col in range(n_per_class):
            ax = axes[row, col]
            if col < len(idxs):
                item = dataset_subset[idxs[col]]
                img  = item[0].permute(1, 2, 0).numpy()
                ax.imshow(np.clip(img, 0, 1))
            else:
                ax.set_facecolor('#111111')
            ax.axis('off')
            if col == 0:
                ax.set_title(str(display_name), fontsize=9, loc='left', pad=4)

    fig.suptitle(title, fontsize=14, fontweight='bold')
    plt.tight_layout()
    plt.show()

base = dog_cat_dataset_train

# ── 1. Label (species) ────────────────────────────────────────────
label_map = {0: 'cat', 1: 'dog'}
show_examples('Examples by Species', 'labels', label_map, train_dataset)

# ── 2. Color ──────────────────────────────────────────────────────
color_map = base.color_to_idx   # {color_str: idx}
show_examples('Examples by Color', 'colors', color_map, train_dataset)

# ── 3. Breed ──────────────────────────────────────────────────────
breed_map = base.breed_to_idx   # {breed_str: idx}
show_examples('Examples by Breed', 'breeds', breed_map, train_dataset)

Lots of organizing info that we can provide! But, how do we do this?

First, we'll create separate embeddings for each category: species, color, and breed.

self.species_embed = nn.Embedding(num_species, 4)    # cat/dog → 4 dims
self.color_embed = nn.Embedding(num_colors, 8)       # 7 colors → 8 dims
self.breed_embed = nn.Embedding(num_breeds, 16)      # ~100 breeds → 16 dims

Each embedding is a learned lookup table that maps a discrete category index to a dense vector. The dimensionality of each embedding should roughly scale with the complexity of what it encodes — species (2 classes) needs very little, breed (~100 classes with subtle distinctions) needs more.

Then, we need to figure out how to introduce the embeddings into the VAE procedure.

Conditioning the Decoder

The decoder needs to know what to generate. We concatenate all condition embeddings into a single conditioning vector and append it to $z$ before decoding:

# In forward():
c_species = self.species_embed(species_label)   # (B, 4)
c_color = self.color_embed(color_label)         # (B, 8)
c_breed = self.breed_embed(breed_label)         # (B, 16)
condition = torch.cat([c_species, c_color, c_breed], dim=1)  # (B, 28)

z_conditioned = torch.cat([z, condition], dim=1)  # (B, latent_dim + 28)
z_projected = self.fc_decode(z_conditioned).view(-1, 512, 4, 4)
reconstruction = self.decoder(z_projected)

The decoder receives a vector that says both "where in latent space" (from $z$) and "what kind of animal" (from the condition). The latent dimensions encode everything else — pose, lighting, background, expression — while the conditioning handles identity.

Conditioning the Encoder

We also want the encoder to know the class, so it can organize the latent space efficiently. Without encoder conditioning, some latent dimensions get wasted encoding "is this a cat or dog" when that information could flow through the embedding instead.

The standard approach is spatial broadcasting: reshape the condition vector to a spatial tensor and concatenate it as extra input channels:

# In forward():
condition = torch.cat([c_species, c_color, c_breed], dim=1)  # (B, 28)
c_spatial = condition.view(B, 28, 1, 1).expand(B, 28, H, W)  # (B, 28, H, W)
x_conditioned = torch.cat([x, c_spatial], dim=1)              # (B, 3 + 28, H, W)
encoded = self.encoder(x_conditioned)  # first conv takes 31 input channels

Every spatial position in the input now carries the same conditioning signal. The encoder can learn to use this to organize the latent space by class — encoding golden retriever features differently from tabby cat features, rather than mixing everything together.

# ════════════════════════════════════════════════════════════════
# FIGURE 2 — CONDITIONAL VAE  (training architecture)
# ════════════════════════════════════════════════════════════════
 
def figure2_conditional_vae():
    fig, ax = plt.subplots(1, 1, figsize=(32, 18))
    ax.set_xlim(0, 32); ax.set_ylim(0, 18)
    ax.axis('off')
    fig.patch.set_facecolor(BG_COLOR); ax.set_facecolor(BG_COLOR)
 
    enc_y, dec_y = 8.6, 2.2
    L = layout(enc_y, dec_y)
 
    ax.text(16.0, 17.55,
            'Conditional VAE — Embedding Integration',
            ha='center', va='center', fontsize=16, color=TITLE_COLOR,
            fontweight='bold', zorder=5)
 
    # ── EMBEDDING BLOCKS (above encoder) ──
    emb_y = enc_y + BH + 1.40
    emb_h = 1.30; emb_w = 2.40; emb_gap = 0.35
    emb_start_x = L['ex'][0] + 0.20
 
    emb_specs = [
        ("Species", "cat / dog --> R⁴", "Embedding(2, 4)"),
        ("Color",   "7 colors --> R⁸",  "Embedding(7, 8)"),
        ("Breed",   "~100 --> R¹⁶",     "Embedding(100, 16)"),
    ]
 
    emb_xs = []
    for i, (name, mapping, impl) in enumerate(emb_specs):
        ex_i = emb_start_x + i * (emb_w + emb_gap)
        emb_xs.append(ex_i)
        # Shadow + box
        ax.add_patch(FancyBboxPatch(
            (ex_i+0.05, emb_y-0.05), emb_w, emb_h,
            boxstyle="round,pad=0.06", facecolor='black',
            edgecolor='none', alpha=0.4, zorder=1))
        ax.add_patch(FancyBboxPatch(
            (ex_i, emb_y), emb_w, emb_h,
            boxstyle="round,pad=0.06", facecolor=COND_COLOR,
            edgecolor=COND_EDGE, linewidth=2.5, zorder=2))
        ax.text(ex_i + emb_w/2, emb_y + emb_h - 0.12, name,
                ha='center', va='top', fontsize=10, color=NOTE_COLOR,
                fontweight='bold', zorder=3)
        ax.text(ex_i + emb_w/2, emb_y + emb_h - 0.42, mapping,
                ha='center', va='top', fontsize=8.5, color=OP_COLOR, zorder=3)
        ax.text(ex_i + emb_w/2, emb_y + emb_h - 0.65, impl,
                ha='center', va='top', fontsize=8, color=GRAY, zorder=3)
 
    # ── CONCAT BOX (combines embeddings) ──
    cat_x = emb_xs[2] + emb_w + 0.50
    cat_y = emb_y + 0.15
    cat_w = 2.00; cat_h = 1.00
 
    ax.add_patch(FancyBboxPatch(
        (cat_x, cat_y), cat_w, cat_h, boxstyle="round,pad=0.06",
        facecolor=COND_COLOR, edgecolor=COND_EDGE,
        linewidth=2.5, zorder=2))
    ax.text(cat_x + cat_w/2, cat_y + cat_h - 0.12, "Concat",
            ha='center', va='top', fontsize=10, color=NOTE_COLOR,
            fontweight='bold', zorder=3)
    ax.text(cat_x + cat_w/2, cat_y + cat_h - 0.42,
            "c = [c_s; c_col; c_br]",
            ha='center', va='top', fontsize=8.5, color=OP_COLOR, zorder=3)
    ax.text(cat_x + cat_w/2, cat_y + cat_h - 0.65, "c ∈ R²⁸",
            ha='center', va='top', fontsize=9, color=DIM_COLOR,
            fontweight='bold', zorder=3)
 
    for ex_i in emb_xs:
        arrow(ax, ex_i + emb_w, emb_y + emb_h/2,
              cat_x, cat_y + cat_h/2, color=COND_EDGE)
 
    # ── ENCODER CONDITIONING ──
    # c is spatially broadcast and concatenated to input channels
    enc_cond_x = cat_x + cat_w/2
    enc_cond_note_y = enc_y + BH + 0.30
 
    note_box(ax, L['ex'][0] - 0.05, enc_cond_note_y, BW + 2.80, 0.55,
             'Input: [x ; broadcast(c)]  ∈ R⁶⁴ˣ⁶⁴ˣ⁽³⁺²⁸⁾',
             fontsize=9, face='#0f3460', edge=COND_EDGE, text_color=DIM_COLOR)
 
    # Arrow from concat down to input label
    arrow_bend(ax, cat_x + cat_w/2, cat_y,
               L['ex'][0] + (BW + 2.80)/2 - 0.05, enc_cond_note_y + 0.55,
               rad=-0.3, color=COND_EDGE)
 
    # Arrow from input label to encoder
    arrow(ax, L['ex'][0] + BW/2, enc_cond_note_y,
          L['ex'][0] + BW/2, enc_y + BH + 0.02, color=ARROW_COLOR)
 
    # Encoder
    # Modify first block to show 31 input channels
    enc_specs_cond = list(ENC_SPECS)
    enc_specs_cond[0] = (
        "Enc Block 1",
        ["Conv 31-->64, s=1", "GN * ReLU",
         "Conv 64-->128, s=2", "GN * ReLU", "ResBlock(128)"],
        "64x64x31", "32x32x128")
    for i, (t, ops, ind, outd) in enumerate(enc_specs_cond):
        draw_block(ax, L['ex'][i], enc_y, BW, BH, t, ops, ind, outd,
                   ENC_COLOR, ENC_EDGE)
    for i in range(3):
        arrow(ax, L['ex'][i]+BW, enc_y+BH/2, L['ex'][i+1], enc_y+BH/2)
 
    # Bottleneck & latent (same as base)
    draw_bottleneck(ax, L, enc_y)
    draw_latent(ax, L)
 
    # ── DECODER CONDITIONING ──
    # Concat c to z before fc_decode
    dec_cat_x = L['lat_x'] + FBW + 0.40
    dec_cat_y = L['lat_y'] + 0.10
    dec_cat_w = 2.40; dec_cat_h = 1.75
 
    ax.add_patch(FancyBboxPatch(
        (dec_cat_x + 0.05, dec_cat_y - 0.05), dec_cat_w, dec_cat_h,
        boxstyle="round,pad=0.06", facecolor='black',
        edgecolor='none', alpha=0.4, zorder=1))
    ax.add_patch(FancyBboxPatch(
        (dec_cat_x, dec_cat_y), dec_cat_w, dec_cat_h,
        boxstyle="round,pad=0.06", facecolor=COND_COLOR,
        edgecolor=GOLD_EDGE, linewidth=2.5, zorder=2))
    ax.text(dec_cat_x + dec_cat_w/2, dec_cat_y + dec_cat_h - 0.12,
            "Concat z + c", ha='center', va='top', fontsize=10,
            color=NOTE_COLOR, fontweight='bold', zorder=3)
    ax.text(dec_cat_x + dec_cat_w/2, dec_cat_y + dec_cat_h - 0.42,
            "[z ; c]", ha='center', va='top', fontsize=9,
            color=OP_COLOR, zorder=3)
    ax.text(dec_cat_x + dec_cat_w/2, dec_cat_y + dec_cat_h - 0.68,
            "z ∈ R²⁵⁶,  c ∈ R²⁸", ha='center', va='top', fontsize=8,
            color=GRAY, zorder=3)
    ax.text(dec_cat_x + dec_cat_w/2, dec_cat_y + dec_cat_h - 0.92,
            "--> R²⁸⁴", ha='center', va='top', fontsize=9,
            color=DIM_COLOR, fontweight='bold', zorder=3)
 
    # z --> concat box
    arrow(ax, L['lat_x'] + FBW, L['lat_y'] + FBH/2,
          dec_cat_x, dec_cat_y + dec_cat_h * 0.55, color=GOLD_EDGE)
 
    # c --> concat box (bend down from concat at top)
    arrow_bend(ax, cat_x + cat_w, cat_y + cat_h/2,
               dec_cat_x + dec_cat_w/2, dec_cat_y + dec_cat_h,
               rad=-0.4, color=COND_EDGE)
 
    # concat box --> fc_decode
    # fc_decode now takes 284 input
    fcd_in = "z+c  (284)"
    draw_block(ax, L['fcd_x'], L['fcd_y'], FBW, FBH, "fc_decode",
               ["Linear(284 --> 8192)", "view(−1, 512, 4, 4)"],
               fcd_in, "4x4x512", FLAT_COLOR, FLAT_EDGE, is_flat=True)
 
    arrow(ax, dec_cat_x + dec_cat_w/2, dec_cat_y,
          L['fcd_x'] + FBW/2, L['fcd_y'] + FBH, color=ARROW_COLOR)
 
    # Decoder
    draw_decoder(ax, L, dec_y)
    note_box(ax, L['dx'][3]-0.05, dec_y-0.85, BW+0.10, 0.55,
             'Reconstruction  x_hat ∈ R⁶⁴ˣ⁶⁴ˣ³', face='#2e0a2e')
    arrow(ax, L['dx'][3]+BW/2, dec_y, L['dx'][3]+BW/2, dec_y-0.28)
 
    # ── LOSS NOTE ──
    lx = L['dx'][3] + BW + 0.50; ly = dec_y - 0.85
    lw, lh = 7.5, 1.20
    dashed_box(ax, lx, ly, lw, lh, edge=MATH_COLOR)
    ax.text(lx + lw/2, ly + lh - 0.10,
            'Loss unchanged — conditioning flows through architecture only',
            ha='center', va='top', fontsize=10, color=MATH_COLOR,
            fontweight='bold', zorder=8)
    ax.text(lx + 0.20, ly + lh - 0.50,
            "L  =  L1(x, x_hat)  +  lam * LPIPS(x, x_hat)  +  D_KL( Q(z|x,c) || N(0,I) )",
            ha='left', va='top', fontsize=9, color=MATH_COLOR,
            family='monospace', zorder=8)
 
    # Section labels
    enc_mid = (L['ex'][0] + L['ex'][3] + BW) / 2
    bot_mid = (L['stack_x'] + L['lat_x'] + FBW) / 2
    for txt, xp, yp in [("◀  ENCODER (conditioned)  ▶", enc_mid, enc_y+BH+1.10),
                         ("◀  BOTTLENECK  ▶", bot_mid, enc_y+BH+1.10),
                         ("◀  DECODER  ▶",
                          (L['dx'][3] + L['dx'][0] + BW)/2, 1.4)]:
        ax.text(xp, yp, txt, ha='center', va='center', fontsize=10,
                color=GRAY, fontstyle='italic', zorder=5)
 
    # Legend
    legend_items = [
        mpatches.Patch(facecolor=ENC_COLOR, edgecolor=ENC_EDGE, label='Encoder Block'),
        mpatches.Patch(facecolor=FLAT_COLOR, edgecolor=FLAT_EDGE, label='Bottleneck (Linear)'),
        mpatches.Patch(facecolor=FLAT_COLOR, edgecolor=GOLD_EDGE, label='Latent Vector z'),
        mpatches.Patch(facecolor=COND_COLOR, edgecolor=COND_EDGE, label='Conditioning'),
        mpatches.Patch(facecolor=DEC_COLOR, edgecolor=DEC_EDGE, label='Decoder Block'),
    ]
    ax.legend(handles=legend_items, loc='lower right', fontsize=9,
              facecolor=BG_COLOR, edgecolor='white',
              labelcolor='white', framealpha=0.85)
 
    plt.tight_layout()
    plt.savefig('fig2_conditional_vae.png', dpi=150, bbox_inches='tight',
                facecolor=fig.get_facecolor())
    plt.show()
    print("Saved: fig2_conditional_vae.png")

figure2_conditional_vae()

Saved: fig2_conditional_vae.png

With this structure in mind, let's create our ConditionalVAE class.

import torch
import torch.nn as nn
import torch.nn.functional as F
import lpips
import numpy as np
from collections import Counter


class ResBlock(nn.Module):
    def __init__(self, channels):
        super().__init__()
        self.block = nn.Sequential(
            nn.Conv2d(channels, channels, kernel_size=3, padding=1),
            nn.GroupNorm(32, channels),
            nn.ReLU(),
            nn.Conv2d(channels, channels, kernel_size=3, padding=1),
            nn.GroupNorm(32, channels),
        )
        self.relu = nn.ReLU()

    def forward(self, x):
        return self.relu(self.block(x) + x)


class EmpiricalConditional:
    """
    Stores the empirical joint distribution P(species, color, breed)
    from training data. Supports sampling with any subset of attributes
    specified, drawing the rest from the conditional distribution.
    """
    def __init__(self, species_list, color_list, breed_list):
        """
        Args:
            species_list: list/array of species indices for all training images
            color_list:   list/array of color indices for all training images
            breed_list:   list/array of breed indices for all training images
        """
        self.combos = list(zip(
            np.asarray(species_list),
            np.asarray(color_list),
            np.asarray(breed_list),
        ))
        self.combo_counts = Counter(
            (int(s), int(c), int(b)) for s, c, b in self.combos
        )
        self.unique_combos = list(self.combo_counts.keys())

    def sample(self, n_samples=1, species=None, color=None, breed=None):
        """
        Sample (species, color, breed) triples from the empirical distribution.
        Any attribute set to None is sampled from P(attr | specified attrs).
        Any attribute set to an int is held fixed.

        Returns:
            species_out, color_out, breed_out: numpy arrays of shape (n_samples,)
        """
        # Filter to matching combinations
        valid = [
            (s, c, b) for (s, c, b) in self.unique_combos
            if (species is None or s == species)
            and (color is None or c == color)
            and (breed is None or b == breed)
        ]

        if len(valid) == 0:
            # Fallback: if the specified combination never appeared,
            # sample from the full joint
            valid = self.unique_combos

        weights = np.array([self.combo_counts[v] for v in valid], dtype=float)
        weights /= weights.sum()

        indices = np.random.choice(len(valid), size=n_samples, p=weights)
        sampled = [valid[i] for i in indices]

        species_out = np.array([s for s, c, b in sampled])
        color_out   = np.array([c for s, c, b in sampled])
        breed_out   = np.array([b for s, c, b in sampled])

        return species_out, color_out, breed_out


class ConditionalPrior(nn.Module):
    """
    Learned conditional prior P(z|c).
    Additive main effects + small interaction residual.
    """
    def __init__(self, num_species, num_colors, num_breeds,
                 latent_dim=256, embed_dim=64, interaction_dim=32):
        super().__init__()

        self.species_embed = nn.Embedding(num_species, embed_dim)
        self.color_embed   = nn.Embedding(num_colors,  embed_dim)
        self.breed_embed   = nn.Embedding(num_breeds,  embed_dim)

        self.interaction = nn.Sequential(
            nn.Linear(embed_dim * 3, interaction_dim),
            nn.ReLU(),
            nn.Linear(interaction_dim, embed_dim),
        )

        self.mu_head     = nn.Linear(embed_dim, latent_dim)
        self.logvar_head = nn.Linear(embed_dim, latent_dim)

    def forward(self, species, color, breed):
        e_s = self.species_embed(species)
        e_c = self.color_embed(color)
        e_b = self.breed_embed(breed)

        h_main = e_s + e_c + e_b
        h_interact = self.interaction(torch.cat([e_s, e_c, e_b], dim=1))
        h = h_main + h_interact

        return self.mu_head(h), self.logvar_head(h)

    def sample(self, species, color, breed):
        mu, logvar = self.forward(species, color, breed)
        std = torch.exp(0.5 * logvar)
        return mu + std * torch.randn_like(std)


class ConditionalVAE(nn.Module):
    def __init__(self, hidden_size=256, num_species=2, num_colors=7,
                 num_breeds=23, empirical_dist=None):
        super(ConditionalVAE, self).__init__()
        self.hidden_size = hidden_size
        self.empirical_dist = empirical_dist

        # ── Conditioning embeddings (for encoder/decoder) ─────────
        self.species_embed = nn.Embedding(num_species, 4)
        self.color_embed   = nn.Embedding(num_colors,  8)
        self.breed_embed   = nn.Embedding(num_breeds,  16)
        self.cond_dim = 4 + 8 + 16

        # ── Learned conditional prior P(z|c) ──────────────────────
        self.prior = ConditionalPrior(
            num_species, num_colors, num_breeds,
            latent_dim=hidden_size,
            embed_dim=64,
            interaction_dim=32,
        )

        # ── Encoder ───────────────────────────────────────────────
        self.encoder = nn.Sequential(
            nn.Conv2d(3 + self.cond_dim, 32, kernel_size=3, stride=1, padding=1),
            nn.GroupNorm(32, 32),
            nn.ReLU(),

            nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1),
            nn.GroupNorm(32, 64),
            nn.ReLU(),
            ResBlock(64),

            nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1),
            nn.GroupNorm(32, 128),
            nn.ReLU(),
            ResBlock(128),

            nn.Conv2d(128, 256, kernel_size=3, stride=2, padding=1),
            nn.GroupNorm(32, 256),
            nn.ReLU(),
            ResBlock(256),

            nn.Conv2d(256, 512, kernel_size=3, stride=2, padding=1),
            nn.GroupNorm(32, 512),
            nn.ReLU(),
            ResBlock(512),

            nn.Flatten(),
        )

        self.fc_mu     = nn.Linear(512 * 4 * 4, hidden_size)
        self.fc_logvar = nn.Linear(512 * 4 * 4, hidden_size)
        self.fc_decode = nn.Linear(hidden_size + self.cond_dim, 512 * 4 * 4)

        # ── Decoder ───────────────────────────────────────────────
        self.decoder = nn.Sequential(
            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),
            nn.Conv2d(512, 256, kernel_size=3, padding=1),
            nn.GroupNorm(32, 256),
            nn.ReLU(),
            ResBlock(256),

            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),
            nn.Conv2d(256, 128, kernel_size=3, padding=1),
            nn.GroupNorm(32, 128),
            nn.ReLU(),
            ResBlock(128),

            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),
            nn.Conv2d(128, 64, kernel_size=3, padding=1),
            nn.GroupNorm(32, 64),
            nn.ReLU(),
            ResBlock(64),

            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),
            nn.Conv2d(64, 32, kernel_size=3, padding=1),
            nn.GroupNorm(32, 32),
            nn.ReLU(),
            ResBlock(32),

            nn.Conv2d(32, 3, kernel_size=1, stride=1),
            nn.Sigmoid(),
        )

    def get_last_decoder_layer(self):
        return self.decoder[-2].weight

    def reparameterize(self, mu, logvar):
        std = torch.exp(0.5 * logvar)
        return mu + std * torch.randn_like(std)

    def encode_condition(self, species, color, breed):
        return torch.cat([
            self.species_embed(species),
            self.color_embed(color),
            self.breed_embed(breed),
        ], dim=1)

    def get_prior(self, species, color, breed):
        return self.prior(species, color, breed)

    def forward(self, x, species, color, breed):
        B, C, H, W = x.shape

        condition = self.encode_condition(species, color, breed)
        c_spatial = condition.view(B, self.cond_dim, 1, 1).expand(B, self.cond_dim, H, W)
        x_cond    = torch.cat([x, c_spatial], dim=1)

        encoded  = self.encoder(x_cond)
        mu_q     = self.fc_mu(encoded)
        logvar_q = torch.clamp(self.fc_logvar(encoded), min=-10, max=10)
        z        = self.reparameterize(mu_q, logvar_q)

        mu_p, logvar_p = self.get_prior(species, color, breed)

        z_cond      = torch.cat([z, condition], dim=1)
        z_projected = self.fc_decode(z_cond).view(-1, 512, 4, 4)
        reconstruction = self.decoder(z_projected)

        return reconstruction, mu_q, logvar_q, mu_p, logvar_p

    @torch.no_grad()
    def generate(self, n_samples=1, species=None, color=None, breed=None):
        """
        Generate images. Any attribute set to None is sampled from the
        empirical conditional distribution P(attr | specified attrs).

        Args:
            n_samples: number of images to generate
            species:   int or None (None = sample from empirical)
            color:     int or None
            breed:     int or None

        Returns:
            generated images tensor (n_samples, 3, 64, 64)
        """
        device = next(self.parameters()).device

        if species is None or color is None or breed is None:
            assert self.empirical_dist is not None, \
                "Must provide empirical_dist to sample unspecified attributes. " \
                "Pass it when constructing the model or set model.empirical_dist."

            sp_np, col_np, br_np = self.empirical_dist.sample(
                n_samples=n_samples,
                species=species,
                color=color,
                breed=breed,
            )
            sp  = torch.tensor(sp_np,  dtype=torch.long, device=device)
            col = torch.tensor(col_np, dtype=torch.long, device=device)
            br  = torch.tensor(br_np,  dtype=torch.long, device=device)
        else:
            sp  = torch.full((n_samples,), species, dtype=torch.long, device=device)
            col = torch.full((n_samples,), color,   dtype=torch.long, device=device)
            br  = torch.full((n_samples,), breed,   dtype=torch.long, device=device)

        # Sample z from learned conditional prior P(z|c)
        z = self.prior.sample(sp, col, br)

        # Build first-stage condition embeddings for decoder
        condition = self.encode_condition(sp, col, br)

        # Decode
        z_cond      = torch.cat([z, condition], dim=1)
        z_projected = self.fc_decode(z_cond).view(-1, 512, 4, 4)
        return self.decoder(z_projected)

    @torch.no_grad()
    def edit(self, x, orig_species, orig_color, orig_breed,
             new_species=None, new_color=None, new_breed=None):
        """
        Encode with original labels, decode with edited labels.
        None = keep the original attribute.
        """
        B = x.shape[0]

        condition_orig = self.encode_condition(orig_species, orig_color, orig_breed)
        c_spatial = condition_orig.view(B, self.cond_dim, 1, 1).expand(B, self.cond_dim, *x.shape[2:])
        x_cond = torch.cat([x, c_spatial], dim=1)

        encoded = self.encoder(x_cond)
        z       = self.fc_mu(encoded)  # posterior mean for clean editing

        edit_species = new_species if new_species is not None else orig_species
        edit_color   = new_color   if new_color   is not None else orig_color
        edit_breed   = new_breed   if new_breed   is not None else orig_breed
        condition_edit = self.encode_condition(edit_species, edit_color, edit_breed)

        z_cond      = torch.cat([z, condition_edit], dim=1)
        z_projected = self.fc_decode(z_cond).view(-1, 512, 4, 4)
        return self.decoder(z_projected)

# ══════════════════════════════════════════════════════════════════
# Loss functions
# ══════════════════════════════════════════════════════════════════

class LPIPSLoss(nn.Module):
    def __init__(self):
        super().__init__()
        self.lpips = lpips.LPIPS(net='vgg')
        for param in self.parameters():
            param.requires_grad = False

    def forward(self, reconstruction, target):
        return self.lpips(reconstruction * 2 - 1, target * 2 - 1).mean()


def compute_adaptive_weight(loss_a, loss_b, last_layer_weight):
    grads_a = torch.autograd.grad(loss_a, last_layer_weight, retain_graph=True)[0]
    grads_b = torch.autograd.grad(loss_b, last_layer_weight, retain_graph=True)[0]
    weight = torch.norm(grads_a) / (torch.norm(grads_b) + 1e-4)
    return torch.clamp(weight, 0.0, 1e4).detach()


def kl_two_gaussians(mu_q, logvar_q, mu_p, logvar_p, free_bits=0.1):
    kl_per_dim = 0.5 * (
        logvar_p - logvar_q - 1
        + (logvar_q.exp() + (mu_q - mu_p).pow(2)) / logvar_p.exp()
    )
    return torch.sum(torch.clamp(kl_per_dim, min=free_bits))


def prior_kl_to_global(mu_p, logvar_p):
    return -0.5 * torch.sum(1 + logvar_p - mu_p.pow(2) - logvar_p.exp())


def conditional_vae_loss(reconstruction, x, mu_q, logvar_q, mu_p, logvar_p,
                         perceptual_loss_fn, last_layer_weight,
                         beta=1.0, gamma=0.1, free_bits=0.1):
    B = x.shape[0]

    nll_loss = F.l1_loss(reconstruction, x, reduction='sum') / B
    p_loss = perceptual_loss_fn(reconstruction, x)
    d_weight = compute_adaptive_weight(nll_loss, p_loss, last_layer_weight)
    wp_loss = d_weight * p_loss

    kl_posterior = kl_two_gaussians(mu_q, logvar_q, mu_p, logvar_p,
                                    free_bits=free_bits) / B
    kl_prior = prior_kl_to_global(mu_p, logvar_p) / B

    total = nll_loss + wp_loss + beta * kl_posterior + gamma * kl_prior

    return (total,
            nll_loss.item(),
            kl_posterior.item(),
            kl_prior.item(),
            wp_loss.item(),
            d_weight.item())

# ══════════════════════════════════════════════════════════════════
# Setup
# ══════════════════════════════════════════════════════════════════

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
hidden_size = 256

# Build empirical distribution from your labels CSV
import pandas as pd

df = pd.read_csv('animal_labels.csv')

# Build label-to-index mappings
species_to_idx = {v: i for i, v in enumerate(sorted(df['label'].unique()))}
color_to_idx   = {v: i for i, v in enumerate(sorted(df['color'].unique()))}
breed_to_idx   = {v: i for i, v in enumerate(sorted(df['breed'].unique()))}

empirical_dist = EmpiricalConditional(
    species_list=df['label'].map(species_to_idx).values,
    color_list=df['color'].map(color_to_idx).values,
    breed_list=df['breed'].map(breed_to_idx).values,
)

model = ConditionalVAE(
    hidden_size=hidden_size,
    num_species=len(species_to_idx),
    num_colors=len(color_to_idx),
    num_breeds=len(breed_to_idx),
    empirical_dist=empirical_dist,
).to(device)

perceptual_loss_fn = LPIPSLoss().to(device)

from torch import optim
optimizer = optim.AdamW(model.parameters(), lr=1e-3)
num_epochs = 500

## KL Annealing

import numpy as np

def get_beta(epoch, warmup_start=10, warmup_end=20, beta_max=1.0):
    if epoch <= warmup_start:
        return 0.0
    elif epoch <= warmup_end:
        t = epoch - warmup_start
        warmup_steps = warmup_end - warmup_start
        return beta_max * 0.5 * (1 - np.cos(np.pi * t / warmup_steps))
    else:
        return beta_max
    
def get_gamma(epoch, hold_epochs=25, anneal_epochs=50, gamma_start=0.1, gamma_end=0.01):
    if epoch <= hold_epochs:
        return gamma_start
    elif epoch <= hold_epochs + anneal_epochs:
        progress = (epoch - hold_epochs) / anneal_epochs
        return gamma_start + (gamma_end - gamma_start) * progress
    else:
        return gamma_end

Setting up [LPIPS] perceptual loss: trunk [vgg], v[0.1], spatial [off]

/home/kmcalist/.local/lib/python3.10/site-packages/torchvision/models/_utils.py:208: UserWarning:

The parameter 'pretrained' is deprecated since 0.13 and may be removed in the future, please use 'weights' instead.

/home/kmcalist/.local/lib/python3.10/site-packages/torchvision/models/_utils.py:223: UserWarning:

Arguments other than a weight enum or `None` for 'weights' are deprecated since 0.13 and may be removed in the future. The current behavior is equivalent to passing `weights=VGG16_Weights.IMAGENET1K_V1`. You can also use `weights=VGG16_Weights.DEFAULT` to get the most up-to-date weights.

Loading model from: /home/kmcalist/.local/lib/python3.10/site-packages/lpips/weights/v0.1/vgg.pth

Since the loss is the same, the only thing that really needs to change is our training loop since we need to unpack conditions. Let's instantiate our model and rebuild our training loop

from tqdm import tqdm
import numpy as np
import torch
from plotly.subplots import make_subplots
import ipywidgets as widgets
from IPython.display import display
import base64, io
from PIL import Image as PILImage

import plotly.graph_objects as go

# =============================================================================
# DASHBOARD SETUP  — Plotly + ipywidgets (survives nbconvert → HTML)
# =============================================================================

history = {
    'epoch': [], 'total': [], 'nll': [], 'kld_posterior': [],
    'kld_prior': [], 'perceptual': [], 'p_weight': [],
}
snapshot_store = {}

# ── Loss figure (6 panels now) ────────────────────────────────────
loss_fig = make_subplots(
    rows=1, cols=6,
    subplot_titles=(
        'Total Loss', 'NLL (Reconstruction)', 'KL (Q‖P(z|c))',
        'KL (P(z|c)‖N(0,I))', 'Perceptual Loss', 'Perceptual Weight',
    ),
)
trace_specs = [
    ('Total',       '#e94560'),
    ('NLL',         '#a0c4ff'),
    ('KL Post',     '#ffd700'),
    ('KL Prior',    '#ff69b4'),
    ('Perceptual',  '#90ee90'),
    ('P-Weight',    '#ff9f40'),
]
for i, (name, color) in enumerate(trace_specs):
    loss_fig.add_trace(go.Scatter(
        x=[], y=[], mode='lines',
        line=dict(color=color, width=2), name=name,
    ), row=1, col=i + 1)

loss_fig.update_layout(
    height=320, paper_bgcolor='#1a1a2e', plot_bgcolor='#0d0d1a',
    font=dict(color='white'), showlegend=False,
    margin=dict(t=40, b=30, l=40, r=20),
)
for i in range(1, 7):
    loss_fig.update_xaxes(gridcolor='#222244', title_text='Epoch', row=1, col=i)
    loss_fig.update_yaxes(gridcolor='#222244', row=1, col=i)

loss_widget = go.FigureWidget(loss_fig)


def _update_loss_plots():
    ep = history['epoch']
    keys = ['total', 'nll', 'kld_posterior', 'kld_prior', 'perceptual', 'p_weight']
    with loss_widget.batch_update():
        for i, k in enumerate(keys):
            loss_widget.data[i].x = ep
            loss_widget.data[i].y = history[k]


# ── Snapshot viewer ───────────────────────────────────────────────
def _tensor_to_b64(img_np):
    arr = (np.clip(img_np, 0, 1) * 255).astype(np.uint8)
    buf = io.BytesIO()
    PILImage.fromarray(arr).save(buf, format='PNG')
    return 'data:image/png;base64,' + base64.b64encode(buf.getvalue()).decode()

ROW_LABELS = ['Train  Orig / Recon', 'Val  Orig / Recon', 'Generated']

img_widgets = [[widgets.Image(format='png', width=80, height=80)
                for _ in range(10)] for _ in range(3)]

row_boxes = []
for r in range(3):
    label = widgets.Label(
        value=ROW_LABELS[r],
        layout=widgets.Layout(width='130px', display='flex',
                              align_items='center'))
    row_box = widgets.HBox(
        [label] + img_widgets[r],
        layout=widgets.Layout(align_items='center'))
    row_boxes.append(row_box)

snap_label = widgets.Label(
    value='Snapshot — no data yet',
    layout=widgets.Layout(margin='4px 0 4px 0'))
snap_label.style.font_size = '13px'

snap_idx = [0]

btn_prev = widgets.Button(description='◀  Prev', button_style='',
                          layout=widgets.Layout(width='100px'))
btn_next = widgets.Button(description='Next  ▶', button_style='',
                          layout=widgets.Layout(width='100px'))
for b in (btn_prev, btn_next):
    b.style.button_color = '#16213e'
    b.style.text_color   = 'white'

epoch_dropdown = widgets.Dropdown(
    options=[], description='Epoch:', layout=widgets.Layout(width='150px'))


def _render_snapshot(epoch):
    imgs = snapshot_store.get(epoch)
    if imgs is None:
        return
    for r in range(3):
        for c in range(10):
            b64 = _tensor_to_b64(imgs[r * 10 + c])
            raw = base64.b64decode(b64.split(',')[1])
            img_widgets[r][c].value = raw
    snap_label.value = f'Snapshot — Epoch {epoch}'


def _on_prev(_b):
    keys = sorted(snapshot_store.keys())
    if not keys: return
    snap_idx[0] = max(0, snap_idx[0] - 1)
    epoch_dropdown.value = keys[snap_idx[0]]


def _on_next(_b):
    keys = sorted(snapshot_store.keys())
    if not keys: return
    snap_idx[0] = min(len(keys) - 1, snap_idx[0] + 1)
    epoch_dropdown.value = keys[snap_idx[0]]


def _on_dropdown_change(change):
    if change['name'] == 'value' and change['new'] is not None:
        keys = sorted(snapshot_store.keys())
        snap_idx[0] = keys.index(change['new'])
        _render_snapshot(change['new'])


btn_prev.on_click(_on_prev)
btn_next.on_click(_on_next)
epoch_dropdown.observe(_on_dropdown_change)

nav_bar     = widgets.HBox([btn_prev, epoch_dropdown, btn_next])
snap_viewer = widgets.VBox([snap_label, nav_bar] + row_boxes)
dashboard   = widgets.VBox([loss_widget, snap_viewer])
display(dashboard)


def _capture_snapshot(epoch):
    model.eval()
    imgs = []
    with torch.no_grad():
        # ── Train: orig + recon (5 pairs) ────────────────────────
        tb, tc_s, tc_b, tc_c = next(iter(train_loader))
        tb   = tb[:5].to(device)
        tc_s = tc_s[:5].to(device)
        tc_b = tc_b[:5].to(device)
        tc_c = tc_c[:5].to(device)
        tr, _, _, _, _ = model(tb, tc_s, tc_c, tc_b)
        for i in range(5):
            imgs.append(tb[i].cpu().permute(1, 2, 0).numpy())
            imgs.append(tr[i].cpu().permute(1, 2, 0).numpy())

        # ── Val: orig + recon (5 pairs) ───────────────────────────
        vb, vc_s, vc_b, vc_c = next(iter(val_loader))
        vb   = vb[:5].to(device)
        vc_s = vc_s[:5].to(device)
        vc_b = vc_b[:5].to(device)
        vc_c = vc_c[:5].to(device)
        vr, _, _, _, _ = model(vb, vc_s, vc_c, vc_b)
        for i in range(5):
            imgs.append(vb[i].cpu().permute(1, 2, 0).numpy())
            imgs.append(vr[i].cpu().permute(1, 2, 0).numpy())

        # ── Generated: sample from learned conditional prior ──────
        g_s, g_c, g_b = model.empirical_dist.sample(n_samples=10)
        g_s = torch.tensor(g_s, dtype=torch.long, device=device)
        g_c = torch.tensor(g_c, dtype=torch.long, device=device)
        g_b = torch.tensor(g_b, dtype=torch.long, device=device)

        # Sample z from P(z|c) via the prior network
        z_gen = model.prior.sample(g_s, g_c, g_b)

        # Build decoder conditioning from first-stage embeddings
        condition = model.encode_condition(g_s, g_c, g_b)
        z_cond = torch.cat([z_gen, condition], dim=1)
        gen = model.decoder(model.fc_decode(z_cond).view(-1, 512, 4, 4))

        for i in range(10):
            imgs.append(gen[i].cpu().permute(1, 2, 0).numpy())

    snapshot_store[epoch] = imgs
    epoch_dropdown.options = sorted(snapshot_store.keys())
    snap_idx[0] = len(snapshot_store) - 1
    epoch_dropdown.value = epoch
    model.train()


# =============================================================================
# TRAINING LOOP
# =============================================================================

best_avg_loss = float('inf')
global_step   = 0

for epoch in range(1, num_epochs + 1):
    model.train()
    train_loss       = 0.0
    total_nll        = 0.0
    total_kld_post   = 0.0
    total_kld_prior  = 0.0
    total_p_loss     = 0.0
    total_p_weight   = 0.0

    beta  = get_beta(epoch, warmup_start=5, warmup_end=15, beta_max=1.0)
    gamma = get_gamma(epoch, hold_epochs=20, anneal_epochs=50, gamma_start=0.1, gamma_end=0.01)

    pbar = tqdm(train_loader,
                desc=f"Epoch {epoch:3d}/{num_epochs}", leave=False)

    for batch_idx, (data, cond_s, cond_b, cond_c) in enumerate(pbar):
        data   = data.to(device)
        cond_s = cond_s.to(device)
        cond_b = cond_b.to(device)
        cond_c = cond_c.to(device)

        optimizer.zero_grad()

        # Forward: returns (recon, mu_q, logvar_q, mu_p, logvar_p)
        recon_batch, mu_q, logvar_q, mu_p, logvar_p = model(
            data, cond_s, cond_c, cond_b
        )

        last_layer_weight = model.get_last_decoder_layer()

        loss, nll, kld_post, kld_prior, p_loss, p_weight = conditional_vae_loss(
            reconstruction     = recon_batch,
            x                  = data,
            mu_q               = mu_q,
            logvar_q           = logvar_q,
            mu_p               = mu_p,
            logvar_p           = logvar_p,
            perceptual_loss_fn = perceptual_loss_fn,
            last_layer_weight  = last_layer_weight,
            beta               = beta,
            gamma              = gamma,
            free_bits          = 0.1,
        )

        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=10.0)
        optimizer.step()

        train_loss      += loss.item()
        total_nll       += nll
        total_kld_post  += kld_post
        total_kld_prior += kld_prior
        total_p_loss    += p_loss
        total_p_weight  += p_weight
        global_step     += 1

        pbar.set_postfix({
            'Loss':    f'{train_loss      / (batch_idx+1):.1f}',
            'NLL':     f'{total_nll       / (batch_idx+1):.1f}',
            'KL_Q':    f'{total_kld_post  / (batch_idx+1):.1f}',
            'KL_P':    f'{total_kld_prior / (batch_idx+1):.2f}',
            'PL':      f'{total_p_loss    / (batch_idx+1):.4f}',
            'PW':      f'{total_p_weight  / (batch_idx+1):.2f}',
        })

    # ── Epoch averages ────────────────────────────────────────────
    n              = len(train_loader)
    avg_loss       = train_loss      / n
    avg_nll        = total_nll       / n
    avg_kld_post   = total_kld_post  / n
    avg_kld_prior  = total_kld_prior / n
    avg_p_loss     = total_p_loss    / n
    avg_p_weight   = total_p_weight  / n

    # ── Checkpoint on best training loss ──────────────────────────
    if avg_loss < best_avg_loss:
        best_avg_loss = avg_loss
        torch.save({
            "epoch":            epoch,
            "global_step":      global_step,
            "model_state_dict": model.state_dict(),
            "opt_state_dict":   optimizer.state_dict(),
            "best_avg_loss":    best_avg_loss,
            "hidden_size":      model.hidden_size,
        }, "conditional_vae_prior_best.pth")

    # ── Live dashboard update (skip first 20 epochs — noisy) ─────
    if epoch >= 20:
        history['epoch']        .append(epoch)
        history['total']        .append(avg_loss)
        history['nll']          .append(avg_nll)
        history['kld_posterior'].append(avg_kld_post)
        history['kld_prior']   .append(avg_kld_prior)
        history['perceptual']  .append(avg_p_loss)
        history['p_weight']    .append(avg_p_weight)
        _update_loss_plots()

    # ── Snapshot every 10 epochs ──────────────────────────────────
    if epoch % 10 == 0:
        _capture_snapshot(epoch)

VBox(children=(FigureWidget({
    'data': [{'line': {'color': '#e94560', 'width': 2},
              'mode': 'l…

---------------------------------------------------------------------------
KeyboardInterrupt                         Traceback (most recent call last)
Cell In[20], line 237
    231 recon_batch, mu_q, logvar_q, mu_p, logvar_p = model(
    232     data, cond_s, cond_c, cond_b
    233 )
    235 last_layer_weight = model.get_last_decoder_layer()
--> 237 loss, nll, kld_post, kld_prior, p_loss, p_weight = conditional_vae_loss(
    238     reconstruction     = recon_batch,
    239     x                  = data,
    240     mu_q               = mu_q,
    241     logvar_q           = logvar_q,
    242     mu_p               = mu_p,
    243     logvar_p           = logvar_p,
    244     perceptual_loss_fn = perceptual_loss_fn,
    245     last_layer_weight  = last_layer_weight,
    246     beta               = beta,
    247     gamma              = gamma,
    248     free_bits          = 0.1,
    249 )
    251 loss.backward()
    252 torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=10.0)

Cell In[18], line 42, in conditional_vae_loss(reconstruction, x, mu_q, logvar_q, mu_p, logvar_p, perceptual_loss_fn, last_layer_weight, beta, gamma, free_bits)
     40 nll_loss = F.l1_loss(reconstruction, x, reduction='sum') / B
     41 p_loss = perceptual_loss_fn(reconstruction, x)
---> 42 d_weight = compute_adaptive_weight(nll_loss, p_loss, last_layer_weight)
     43 wp_loss = d_weight * p_loss
     45 kl_posterior = kl_two_gaussians(mu_q, logvar_q, mu_p, logvar_p,
     46                                 free_bits=free_bits) / B

Cell In[18], line 18, in compute_adaptive_weight(loss_a, loss_b, last_layer_weight)
     16 def compute_adaptive_weight(loss_a, loss_b, last_layer_weight):
     17     grads_a = torch.autograd.grad(loss_a, last_layer_weight, retain_graph=True)[0]
---> 18     grads_b = torch.autograd.grad(loss_b, last_layer_weight, retain_graph=True)[0]
     19     weight = torch.norm(grads_a) / (torch.norm(grads_b) + 1e-4)
     20     return torch.clamp(weight, 0.0, 1e4).detach()

File ~/.local/lib/python3.10/site-packages/torch/autograd/__init__.py:515, in grad(outputs, inputs, grad_outputs, retain_graph, create_graph, only_inputs, allow_unused, is_grads_batched, materialize_grads)
    511     result = _vmap_internals._vmap(vjp, 0, 0, allow_none_pass_through=True)(
    512         grad_outputs_
    513     )
    514 else:
--> 515     result = _engine_run_backward(
    516         outputs,
    517         grad_outputs_,
    518         retain_graph,
    519         create_graph,
    520         inputs,
    521         allow_unused,
    522         accumulate_grad=False,
    523     )
    524 if materialize_grads:
    525     if any(
    526         result[i] is None and not is_tensor_like(inputs[i])
    527         for i in range(len(inputs))
    528     ):

File ~/.local/lib/python3.10/site-packages/torch/autograd/graph.py:865, in _engine_run_backward(t_outputs, *args, **kwargs)
    863     unregister_hooks = _register_logging_hooks_on_whole_graph(t_outputs)
    864 try:
--> 865     return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
    866         t_outputs, *args, **kwargs
    867     )  # Calls into the C++ engine to run the backward pass
    868 finally:
    869     if attach_logging_hooks:

KeyboardInterrupt: 

Flexible Generation And Editing: Conditional, Unconditional, and Partial

import torch
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.gridspec import GridSpec

# Assumes: model, device, empirical_dist, species_to_idx, color_to_idx, breed_to_idx
# are already defined. Build reverse lookup dicts for labeling.

idx_to_species = {v: k for k, v in species_to_idx.items()}
idx_to_color   = {v: k for k, v in color_to_idx.items()}
idx_to_breed   = {v: k for k, v in breed_to_idx.items()}

model.eval()


# ═════════════════════════════════════════════════════════════════
# FIGURE 1: 100 UNCONDITIONAL DRAWS (10×10 GRID)
# ═════════════════════════════════════════════════════════════════

with torch.no_grad():
    images_uncond = model.generate(n_samples=100)  # all attributes sampled from joint

fig1, axes1 = plt.subplots(10, 10, figsize=(15, 15), facecolor='white')
fig1.suptitle('100 Unconditional Generations\nAll attributes sampled from P(species, color, breed)',
              fontsize=14, fontweight='bold', y=1.02)

for i in range(10):
    for j in range(10):
        idx = i * 10 + j
        img = images_uncond[idx].cpu().permute(1, 2, 0).numpy()
        axes1[i][j].imshow(np.clip(img, 0, 1))
        axes1[i][j].axis('off')

plt.tight_layout()
plt.savefig('unconditional_100.png', dpi=150, bbox_inches='tight', facecolor='white')
plt.show()

# ═════════════════════════════════════════════════════════════════
# FIGURE 2: TOP 10 MOST COMMON COMBOS (10 rows × 5 images)
# ═════════════════════════════════════════════════════════════════

top10 = empirical_dist.combo_counts.most_common(10)

fig2 = plt.figure(figsize=(10, 18), facecolor='white')
fig2.suptitle('Top 10 Most Common (Species, Color, Breed) Combinations\n5 samples each',
              fontsize=13, fontweight='bold', y=1.01)

# Use GridSpec to give the label column fixed width
gs2 = fig2.add_gridspec(10, 6, width_ratios=[2.5, 1, 1, 1, 1, 1],
                         hspace=0.3, wspace=0.05)

with torch.no_grad():
    for row, ((sp, col, br), count) in enumerate(top10):
        images = model.generate(
            n_samples=5,
            species=sp,
            color=col,
            breed=br,
        )

        sp_name  = idx_to_species[sp]
        col_name = idx_to_color[col]
        br_name  = idx_to_breed[br]

        # Label cell (first column)
        ax_label = fig2.add_subplot(gs2[row, 0])
        ax_label.axis('off')
        ax_label.text(0.95, 0.5,
                      f'{sp_name}\n{col_name}\n{br_name}\n(n={count})',
                      ha='right', va='center', fontsize=9,
                      fontweight='bold', color='#222244',
                      transform=ax_label.transAxes)

        # Image cells (columns 1-5)
        for c in range(5):
            ax_img = fig2.add_subplot(gs2[row, c + 1])
            img = images[c].cpu().permute(1, 2, 0).numpy()
            ax_img.imshow(np.clip(img, 0, 1))
            ax_img.axis('off')

plt.savefig('top10_combos.png', dpi=150, bbox_inches='tight', facecolor='white')
plt.show()

# ═════════════════════════════════════════════════════════════════
# FIGURE 3: EVERY COLOR × {CAT, DOG}, BREED UNCONDITIONED
# ═════════════════════════════════════════════════════════════════

colors = sorted(color_to_idx.items(), key=lambda x: x[1])
n_colors = len(colors)
n_rows = n_colors * 2

fig3 = plt.figure(figsize=(10, n_rows * 1.6 + 1), facecolor='white')
fig3.suptitle('Every Color × Species (breed unconditioned)\n5 samples each',
              fontsize=13, fontweight='bold', y=1.01)

gs3 = fig3.add_gridspec(n_rows, 6, width_ratios=[2.5, 1, 1, 1, 1, 1],
                         hspace=0.3, wspace=0.05)

with torch.no_grad():
    row = 0
    for color_name, color_idx in colors:
        for species_name in ['cat', 'dog']:
            species_idx = species_to_idx[species_name]

            images = model.generate(
                n_samples=5,
                species=species_idx,
                color=color_idx,
                breed=None,
            )

            # Label cell
            ax_label = fig3.add_subplot(gs3[row, 0])
            ax_label.axis('off')
            ax_label.text(0.95, 0.5,
                          f'{species_name}\n{color_name}',
                          ha='right', va='center', fontsize=9,
                          fontweight='bold', color='#222244',
                          transform=ax_label.transAxes)

            # Image cells
            for c in range(5):
                ax_img = fig3.add_subplot(gs3[row, c + 1])
                img = images[c].cpu().permute(1, 2, 0).numpy()
                ax_img.imshow(np.clip(img, 0, 1))
                ax_img.axis('off')

            row += 1

plt.savefig('color_species_grid.png', dpi=150, bbox_inches='tight', facecolor='white')
plt.show()

import torch
import numpy as np
import matplotlib.pyplot as plt

model.eval()

def get_indices_by_attr(dataset_subset, attr, value, limit=10):
    results = []
    base_dataset = dataset_subset.dataset
    for subset_idx, global_idx in enumerate(dataset_subset.indices):
        if getattr(base_dataset, attr)[global_idx] == value:
            results.append(subset_idx)
        if len(results) >= limit:
            break
    return results

# ═════════════════════════════════════════════════════════════════
# FIND 10 GOLDEN RETRIEVERS FROM THE DATASET
# ═════════════════════════════════════════════════════════════════

golden_retriever_idx = dog_cat_dataset_train.breed_to_idx['Labrador retriever']
black_color_idx      = dog_cat_dataset_train.color_to_idx['white']
cat_idx              = dog_cat_dataset_train.label_names_inv = {v: k for k, v in dog_cat_dataset_train.label_names.items()}
cat_idx              = 0  # cat=0 from DogCatDataset

golden_indices = get_indices_by_attr(train_dataset, 'breeds', golden_retriever_idx, limit=10)

golden_images  = []
golden_species = []
golden_colors  = []
golden_breeds  = []

for i in golden_indices:
    img, sp, br, col = train_dataset[i]
    golden_images.append(img)
    golden_species.append(sp)
    golden_colors.append(col)
    golden_breeds.append(br)

print(f"Found {len(golden_images)} Labrador retrievers")

Found 10 Labrador retrievers

dog_cat_dataset_train.color_to_idx

{'black': 0,
 'brown': 1,
 'cream': 2,
 'ginger': 3,
 'golden': 4,
 'gray': 5,
 'white': 6}

img_batch = torch.stack(golden_images).to(device)
sp_batch  = torch.tensor(golden_species, dtype=torch.long, device=device)
col_batch = torch.tensor(golden_colors, dtype=torch.long, device=device)
br_batch  = torch.tensor(golden_breeds, dtype=torch.long, device=device)

# ═════════════════════════════════════════════════════════════════
# ROW 2: RAW RECONSTRUCTION (encode with original labels, decode with original labels)
# ROW 3: COLOR EDIT (golden → black)
# ROW 4: SPECIES EDIT (dog → cat)
# ═════════════════════════════════════════════════════════════════

with torch.no_grad():
    # Reconstruction: edit with no changes
    reconstructions = model.edit(
        img_batch,
        orig_species=sp_batch,
        orig_color=col_batch,
        orig_breed=br_batch,
    )

    # Color edit: golden → black
    black_edits = model.edit(
        img_batch,
        orig_species=sp_batch,
        orig_color=col_batch,
        orig_breed=br_batch,
        new_color=torch.full_like(col_batch, black_color_idx),
    )

    # Species edit: dog → cat
    cat_edits = model.edit(
        img_batch,
        orig_species=sp_batch,
        orig_color=col_batch,
        orig_breed=br_batch,
        new_species=torch.full_like(sp_batch, cat_idx),
    )

# ═════════════════════════════════════════════════════════════════
# 4 × 10 GRID
# ═════════════════════════════════════════════════════════════════

row_data = [
    (img_batch,       'Original'),
    (reconstructions, 'Reconstruction'),
    (black_edits,     'Edit: Color → Black'),
    (cat_edits,       'Edit: Species → Cat'),
]

fig, axes = plt.subplots(4, 10, figsize=(18, 7.5), facecolor='white')
fig.suptitle('Conditional Editing: 10 Golden Retrievers\n'
             'Same z (pose, background, lighting) — only conditioning embeddings change',
             fontsize=13, fontweight='bold', y=1.03)

for row, (images, label) in enumerate(row_data):
    for col in range(10):
        img = images[col].cpu().permute(1, 2, 0).numpy()
        axes[row][col].imshow(np.clip(img, 0, 1))
        axes[row][col].axis('off')

    axes[row][0].set_ylabel(label, fontsize=10, fontweight='bold',
                             rotation=0, labelpad=100, va='center')

plt.tight_layout()
plt.savefig('golden_retriever_edits_4x10.png', dpi=150,
            bbox_inches='tight', facecolor='white')
plt.show()

golden_images

[]