Updated vcae experiment (#45)

* latest experiment

* refactor

* add SVHN

* proper flattening

* increase number of epochs

* add notebooks

* add the option to save results

* add the option to save results+1

* cleanup

* refactor and simplify

Author: tvatter

.gitignore

@@ -11,8 +11,9 @@ MNIST/
 logs/
 examples/MNIST
 examples/logs
+examples/checkpoints
 examples/**/*.gz
-examples/**/*.ipynb
+# examples/**/*.ipynb
 
 # Python-generated files
 __pycache__/
@@ -239,3 +240,6 @@ Temporary Items
 
 # Built Visual Studio Code Extensions
 *.vsix
+
+*.csv
+*.mat

README.md

@@ -41,7 +41,7 @@ Visit the [`./examples/`](https://github.com/TY-Cheng/torchvinecopulib/tree/main
 
 ## Installation
 
-- By `pip` from [`PyPI`](https://pypi.org/project/torchvinecopulib/):
+- By `pip` from [`PyPI`](https://pypi.org/project/torchvinecopulib/) (see the dependencies and uv sections below for CUDA support):
 
 ```bash
 pip install torchvinecopulib torch
@@ -62,10 +62,19 @@ pip install ./dist/torchvinecopulib-1.1.0.tar.gz
 After `git clone https://github.com/TY-Cheng/torchvinecopulib.git`, `cd` into the project root where [`pyproject.toml`](https://github.com/TY-Cheng/torchvinecopulib/blob/main/pyproject.toml) exists,
 
 ```bash
-# inside project root folder
-uv sync --extra cpu -U
-# or
-uv sync --extra cu126 -U
+# From inside the project root folder
+# Create and activate local virtual environment
+uv venv .venv
+source .venv/bin/activate
+
+# Sync dependencies with CPU support (default)
+uv sync --extra cpu
+
+# Or for CUDA 12.6 or 12.8 support (depends on your CUDA version)
+uv sync --extra cu126
+
+# Additionally, to install additional dependencies for the examples
+uv sync --extra examples
 ```
 
 ## Dependencies
@@ -81,6 +90,7 @@ scipy = "*"
 torch = [
     { index = "torch-cpu", extra = "cpu" },
     { index = "torch-cu126", extra = "cu126" },
+    { index = "torch-cu128", extra = "cu128" },
 ]
 ```
 

examples/vcae_mnist/run_seeds.py renamed to examples/vcae/run_seeds.py

@@ -2,17 +2,27 @@
 import logging
 import os
 import sys
+from typing import Union
 
 import pandas as pd
 from tqdm import tqdm
-from vcae_mnist.config import config
-from vcae_mnist.experiment import run_experiment
+from vcae.config import config_mnist, config_svhn
+from vcae.experiment import run_experiment
 
+dataset = "MNIST"  # or "SVHN"
 start = int(sys.argv[1])
 end = int(sys.argv[2])
 
+if dataset == "MNIST":
+    config = config_mnist
+elif dataset == "SVHN":
+    config = config_svhn
+else:
+    raise ValueError(f"Unsupported dataset: {dataset}")
+
+
 # Redirect tqdm and errors to log file
-log_path = f"progress_{start}_{end}.log"
+log_path = f"progress_{dataset}_{start}_{end}.log"
 log_file = open(log_path, "w")
 
 logging.basicConfig(
@@ -33,14 +43,12 @@ def suppress_output():
             logging.getLogger().setLevel(logging_level)
 
 
-results = []
-
-output_path = f"results_{start}_{end}.csv"
-
+results: list[dict[str, Union[float, int, str]]] = []
+output_path = f"results_{dataset}_{start}_{end}.csv"
 for seed in tqdm(range(start, end), desc=f"Seeds {start}-{end}", file=log_file):
     try:
         with suppress_output():
-            result = run_experiment(seed, config)
+            result = run_experiment(seed, config, dataset=dataset)
         df = pd.DataFrame([result])
 
         # Write headers only once

examples/vcae/run_seeds.sh

@@ -0,0 +1,61 @@
+#!/bin/bash
+
+PYTHON_BIN=$(which python)
+USE_NOHUP=false
+
+# Defaults
+START=0
+END=30
+STEP=10
+
+# Parse arguments
+POSITIONAL=()
+while [[ $# -gt 0 ]]; do
+    case "$1" in
+        --nohup)
+            USE_NOHUP=true
+            shift
+            ;;
+        *)
+            POSITIONAL+=("$1")
+            shift
+            ;;
+    esac
+done
+
+# Restore positional args
+set -- "${POSITIONAL[@]}"
+
+# Assign range values if provided
+if [[ $# -ge 1 ]]; then START=$1; fi
+if [[ $# -ge 2 ]]; then END=$2; fi
+if [[ $# -ge 3 ]]; then STEP=$3; fi
+
+# Validate input
+if (( STEP <= 0 )); then
+    echo "Error: STEP must be a positive integer." >&2
+    exit 1
+fi
+
+if (( (END - START) % STEP != 0 )); then
+    echo "Error: (END - START) must be divisible by STEP." >&2
+    exit 1
+fi
+
+echo "Using Python binary: $PYTHON_BIN"
+echo "Using nohup: $USE_NOHUP"
+echo "Range: $START to $END with step $STEP"
+
+# Launch loop
+for ((i = START; i < END; i += STEP)); do
+    j=$((i + STEP))
+    name="seeds_${i}_${j}"
+    echo "Launching $name"
+    if $USE_NOHUP; then
+        nohup "$PYTHON_BIN" run_seeds.py $i $j > logs/$name.log 2>&1 &
+    else
+        "$PYTHON_BIN" run_seeds.py $i $j > logs/$name.log 2>&1 &
+    fi
+done
+
+wait

examples/vcae_mnist/vcae_mnist/__init__.py renamed to examples/vcae/vcae/__init__.py

@@ -0,0 +1,49 @@
+import os
+from dataclasses import dataclass
+
+import torch
+
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+torch.set_float32_matmul_precision("medium")
+
+
+@dataclass
+class Config:
+    # Reproducibility
+    seed: int = 42
+
+    # Training-related
+    data_dir: str = os.environ.get("PATH_DATASETS", ".")
+    save_dir: str = "logs/"
+    batch_size: int = 512 if torch.cuda.is_available() else 64
+    max_epochs: int = 10
+    accelerator: str = DEVICE
+    devices: int = 1
+    num_workers: int = 1  # or min(15, os.cpu_count())
+
+    # Data-related
+    dims: tuple[int, ...] = (1, 28, 28)
+    val_train_split: float = 0.1
+
+    # Model-related
+    hidden_size: int = 64
+    latent_size: int = 10
+    learning_rate: float = 2e-4
+    vine_lambda: float = 0.0
+    # use_mmd: bool = False
+    # mmd_sigmas: list[float] = [1e-1, 1, 10]
+    # mmd_lambda: float = 10.0
+
+config_mnist = Config(
+    max_epochs=10,
+    dims=(1, 28, 28),
+    hidden_size=64,
+    latent_size=10,
+)
+
+config_svhn = Config(
+    max_epochs=50,
+    dims=(3, 32, 32),
+    hidden_size=128,
+    latent_size=32,
+)

examples/vcae/vcae/config.py

@@ -1,13 +1,13 @@
-import copy
 import random
+from typing import Union
 
 import numpy as np
 import pytorch_lightning as pl
 import torch
 
 from .config import DEVICE, Config
 from .metrics import compute_score
-from .model import LitMNISTAutoencoder
+from .model import LitAutoencoder, LitMNISTAutoencoder, LitSVHNAutoencoder
 
 
 def set_seed(seed: int):
@@ -20,11 +20,21 @@ def set_seed(seed: int):
     torch.backends.cudnn.deterministic = True
 
 
-def run_experiment(seed: int, config: Config):
+def run_experiment(
+    seed: int, config: Config, vine_lambda: float = 1.0, dataset: str = "MNIST"
+) -> dict[str, Union[float, int, str]]:
+    # Set the seed for reproducibility
     set_seed(seed)
+    config.seed = seed
 
     # Instantiate the model
-    model_initial = LitMNISTAutoencoder()
+    model_initial: LitAutoencoder
+    if dataset == "MNIST":
+        model_initial = LitMNISTAutoencoder(config)
+    elif dataset == "SVHN":
+        model_initial = LitSVHNAutoencoder(config)
+    else:
+        raise ValueError(f"Unsupported dataset: {dataset}")
 
     # Set up trainer
     trainer_initial = pl.Trainer(
@@ -46,10 +56,16 @@ def run_experiment(seed: int, config: Config):
     model_initial.learn_vine(n_samples=5000)
 
     # Extract test data
-    rep_initial, _, data_initial, _, samples_initial = model_initial.get_data(stage="test")
+    rep_initial, _, data_initial, decoded_initial, samples_initial = model_initial.get_data(
+        stage="test"
+    )
+
+    # Reset the seed for refitting to avoid data leakage
+    set_seed(seed)
 
-    # Deepcopy for refit
-    model_refit = copy.deepcopy(model_initial)
+    # Create a new model with the same configuration but reset vine lambda
+    config.vine_lambda = vine_lambda
+    model_refit = model_initial.copy_with_config(config)
 
     # Set up trainer for refitting
     trainer_refit = pl.Trainer(
@@ -68,21 +84,29 @@ def run_experiment(seed: int, config: Config):
     model_refit.to(DEVICE)
 
     # Extract test data
-    rep_refit, _, data_refit, _, samples_refit = model_refit.get_data(stage="test")
+    rep_refit, _, data_refit, decoded_refit, samples_refit = model_refit.get_data(stage="test")
 
+    assert model_initial.vine is not None
+    assert model_refit.vine is not None
     loglik_initial = model_initial.vine.log_pdf(rep_initial).mean().item()
     loglik_refit = model_refit.vine.log_pdf(rep_refit).mean().item()
 
+    mse_initial = torch.nn.functional.mse_loss(decoded_initial, data_initial).item()
+    mse_refit = torch.nn.functional.mse_loss(decoded_refit, data_refit).item()
+
     sigmas = [1e-3, 1e-2, 1e-1, 1, 10, 100]
-    score_initial = compute_score(data_initial, samples_initial, DEVICE, sigmas=sigmas)
-    score_refit = compute_score(data_refit, samples_refit, DEVICE, sigmas=sigmas)
+    score_initial = compute_score(data_initial, samples_initial, sigmas=sigmas)
+    score_refit = compute_score(data_refit, samples_refit, sigmas=sigmas)
 
     return {
         "seed": seed,
-        "loglik": loglik_initial,
+        "dataset": dataset,
+        "mse_initial": mse_initial,
+        "mse_refit": mse_refit,
+        "loglik_initial": loglik_initial,
         "loglik_refit": loglik_refit,
-        "mmd": score_initial.mmd,
+        "mmd_initial": score_initial.mmd,
         "mmd_refit": score_refit.mmd,
-        "fid": score_initial.fid,
+        "fid_initial": score_initial.fid,
         "fid_refit": score_refit.fid,
     }

examples/vcae_mnist/vcae_mnist/experiment.py renamed to examples/vcae/vcae/experiment.py

@@ -0,0 +1,79 @@
+import numpy as np
+import torch
+from scipy import linalg
+
+
+def mmd(real: torch.Tensor, fake: torch.Tensor, sigmas=[1e-3, 1e-2, 1e-1, 1, 10, 100]):
+    """
+    Differentiable MMD loss using Gaussian kernels with fixed sigmas and
+    distance normalization via Mxx.mean().
+
+    Parameters
+    ----------
+    real : (n, d) tensor
+        Batch of real samples (features or images).
+    fake : (m, d) tensor
+        Batch of generated samples.
+    sigmas : list of float
+        Bandwidths for the RBF kernel. Defaults to wide, fixed list.
+
+    Returns
+    -------
+    mmd : scalar tensor
+        Differentiable scalar loss value.
+    """
+    real = real.view(real.size(0), -1)
+    fake = fake.view(fake.size(0), -1)
+
+    def pairwise_squared_distances(x, y):
+        x_norm = (x**2).sum(dim=1, keepdim=True)
+        y_norm = (y**2).sum(dim=1, keepdim=True)
+        return x_norm + y_norm.T - 2.0 * x @ y.T
+
+    Mxx = pairwise_squared_distances(real, real)
+    Mxy = pairwise_squared_distances(real, fake)
+    Myy = pairwise_squared_distances(fake, fake)
+
+    # Normalization factor based on real-real distances
+    scale = Mxx.mean().detach()
+
+    mmd_total = 0.0
+    for sigma in sigmas:
+        denom = scale * 2.0 * sigma**2
+        Kxx = torch.exp(-Mxx / denom)
+        Kxy = torch.exp(-Mxy / denom)
+        Kyy = torch.exp(-Myy / denom)
+
+        mmd_total += Kxx.mean() + Kyy.mean() - 2.0 * Kxy.mean()
+
+    return mmd_total / len(sigmas)
+
+
+def fid(X, Y):
+    m = X.mean(0)
+    m_w = Y.mean(0)
+    X_np = X.numpy()
+    Y_np = Y.numpy()
+
+    C = np.cov(X_np.transpose())
+    C_w = np.cov(Y_np.transpose())
+    C_C_w_sqrt = linalg.sqrtm(C.dot(C_w), True).real
+
+    score = m.dot(m) + m_w.dot(m_w) - 2 * m_w.dot(m) + np.trace(C + C_w - 2 * C_C_w_sqrt)
+    return np.sqrt(score)
+
+
+class Score:
+    mmd = 0
+    fid = 0
+
+
+def compute_score(real, fake, sigmas=[1e-3, 1e-2, 1e-1, 1, 10, 100]):
+    real = real.to("cpu")
+    fake = fake.to("cpu")
+
+    s = Score()
+    s.mmd = np.sqrt(mmd(real, fake, sigmas).numpy())
+    s.fid = fid(fake, real).numpy()
+
+    return s