vae.py

# -*- coding: utf-8 -*-
"""VAE.ipynb
"""

# ============================
# 1. Imports and Device Setup
# ============================
import os
import csv
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from torch.nn.utils import clip_grad_norm_
from tqdm import tqdm
import matplotlib.pyplot as plt
from pytorch_msssim import ssim  # pip install pytorch-msssim

# Set device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

class ReconstructionDataset(Dataset):
    def __init__(self, X, in_channels=3, transform=None):
        """
        Args:
            X (np.array): Array of images with shape (N, H, W) or (N, H, W, C).
            in_channels (int): Expected number of channels in each image.
            transform: Optional transform to apply to each image tensor.
        """
        self.X = X
        self.in_channels = in_channels
        self.transform = transform

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

    def __getitem__(self, idx):
        # Get image; expected shape: (H, W) or (H, W, C)
        img = self.X[idx]
        # If image is 2D, add a channel dimension.
        if img.ndim == 2:
            img = img[..., None]
        # Check that the image has the expected number of channels.
        if img.shape[-1] != self.in_channels:
            raise ValueError(f"Mismatch canali: image shape {img.shape}, expected in_channels={self.in_channels}.")
        # Convert image to PyTorch tensor and rearrange dimensions to (C, H, W)
        img_t = torch.from_numpy(img).permute(2, 0, 1).float()
        if self.transform:
            img_t = self.transform(img_t)
        # For reconstruction, the target is the same as the input.
        return img_t, img_t

dataset_name = 'pneumonia'
in_channels = 1
pxl = 128

# Load the image data (we ignore labels for reconstruction)
X_train = np.load(f"/content/gdrive/My Drive/Colab Notebooks/MedMnist/dataset/{dataset_name}_{pxl}/X_train.npy")
X_test  = np.load(f"/content/gdrive/My Drive/Colab Notebooks/MedMnist/dataset/{dataset_name}_{pxl}/X_test.npy")

# Normalize to [0,1] (if not already done)
X_train = X_train.astype("float32") / 255.0
X_test  = X_test.astype("float32") / 255.0

print("Dopo conversione, shape X_train:", X_train.shape)
print("Dopo conversione, shape X_test :", X_test.shape)

recon_transform = None

# Create the dataset objects
train_dataset = ReconstructionDataset(X_train, in_channels=in_channels, transform=recon_transform)
test_dataset = ReconstructionDataset(X_test, in_channels=in_channels, transform=recon_transform)

batch_size = 256
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=4)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=4)

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

class BetaVAE(nn.Module):
    def __init__(self, in_channels=3, img_size=128, latent_dim=128, beta=4):
        """
        Args:
            in_channels (int): Number of channels in input images.
            img_size (int): Spatial resolution (assumed square) of the input images.
            latent_dim (int): Dimensionality of the latent space.
            beta (float): Weight for the KL divergence term.

        Note: This architecture assumes the input image size is divisible by 8.
        """
        super(BetaVAE, self).__init__()
        self.beta = beta
        self.in_channels = in_channels
        self.img_size = img_size

        # ----------------------------
        # Encoder
        # ----------------------------
        # We use three convolutional layers with kernel=4, stride=2, padding=1.
        # For an input of size img_size, the output size of each layer is computed as:
        #   out = floor((W - 4 + 2*1)/2) + 1 = floor((W - 2)/2) + 1
        # Thus, after three layers, the spatial dimensions become:
        #   H_enc = img_size_out = (img_size/2^3) if img_size is divisible by 8.
        self.encoder = nn.Sequential(
            nn.Conv2d(in_channels, 32, kernel_size=4, stride=2, padding=1),  # -> (32, img_size/2, img_size/2)
            nn.BatchNorm2d(32),
            nn.ReLU(inplace=True),
            nn.Conv2d(32, 64, kernel_size=4, stride=2, padding=1),           # -> (64, img_size/4, img_size/4)
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.Conv2d(64, 128, kernel_size=4, stride=2, padding=1),          # -> (128, img_size/8, img_size/8)
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True)
        )

        # Compute encoder output shape dynamically.
        dummy = torch.zeros(1, in_channels, img_size, img_size)
        enc_out = self.encoder(dummy)   # shape: (1, 128, H_enc, W_enc)
        self.enc_out_shape = enc_out.shape[2:]  # (H_enc, W_enc)
        self.enc_flat_dim = int(enc_out.view(1, -1).size(1))  # e.g., 128 * (img_size/8)^2

        self.flatten = nn.Flatten()
        self.fc_mu = nn.Linear(self.enc_flat_dim, latent_dim)
        self.fc_logvar = nn.Linear(self.enc_flat_dim, latent_dim)

        # ----------------------------
        # Decoder
        # ----------------------------
        # The decoder will map the latent vector to a feature map with the same
        # flattened dimension as the encoder output, then use transposed convolutions
        # to recover the original image size.
        self.fc_dec = nn.Linear(latent_dim, self.enc_flat_dim)

        # The decoder is designed to mirror the encoder. Starting from a feature map of shape
        # (128, H_enc, W_enc) (with H_enc = img_size/8), we use three ConvTranspose2d layers
        # (each with stride=2, kernel=4, padding=1) to double the spatial dimensions at each step.
        self.decoder = nn.Sequential(
            nn.ConvTranspose2d(128, 64, kernel_size=4, stride=2, padding=1),   # -> (64, img_size/4, img_size/4)
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.ConvTranspose2d(64, 32, kernel_size=4, stride=2, padding=1),    # -> (32, img_size/2, img_size/2)
            nn.BatchNorm2d(32),
            nn.ReLU(inplace=True),
            nn.ConvTranspose2d(32, in_channels, kernel_size=4, stride=2, padding=1),  # -> (in_channels, img_size, img_size)
            nn.Sigmoid()  # Outputs in [0,1]
        )

    def encode(self, x):
        h = self.encoder(x)
        h_flat = self.flatten(h)
        mu = self.fc_mu(h_flat)
        logvar = self.fc_logvar(h_flat)
        return mu, logvar

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

    def decode(self, z):
        h = self.fc_dec(z)
        h = h.view(-1, 128, self.enc_out_shape[0], self.enc_out_shape[1])
        x_recon = self.decoder(h)
        return x_recon

    def forward(self, x):
        mu, logvar = self.encode(x)
        z = self.reparameterize(mu, logvar)
        x_recon = self.decode(z)
        return x_recon, mu, logvar

# Set a beta value and create a helper function.
beta = 0.001  # or change to desired value (e.g., 0.1, 4, etc.)
def create_vae():
    # Pass in_channels and img_size so that the model adapts to any image size.
    # For example, for 128x128 images:
    return BetaVAE(in_channels=in_channels, img_size=pxl, latent_dim=64, beta=beta)

# ============================
# 4. Loss Function and Metrics for VAE
###########################################
def vae_loss(x, x_recon, mu, logvar, beta):
    # Reconstruction loss (MSE, summed over pixels)
    recon_loss = F.mse_loss(x_recon, x, reduction='sum')
    # KL divergence loss
    kl_loss = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
    return recon_loss + beta * kl_loss

def compute_recon_mse(model, dataloader):
    model.eval()
    total_loss = 0.0
    total = 0
    with torch.no_grad():
        for images, _ in dataloader:
            images = images.to(device)
            x_recon, mu, logvar = model(images)
            loss = F.mse_loss(x_recon, images, reduction='sum')
            total_loss += loss.item()
            total += images.size(0)
    return total_loss / total

def compute_ssim_metric(model, dataloader):
    model.eval()
    ssim_total = 0.0
    count = 0
    with torch.no_grad():
        for images, _ in dataloader:
            images = images.to(device)
            x_recon, mu, logvar = model(images)
            batch_ssim = ssim(x_recon, images, data_range=1.0, size_average=True)
            ssim_total += batch_ssim.item() * images.size(0)
            count += images.size(0)
    return ssim_total / count

# ============================
# 5. Define run_experiment for β-VAE Training
# ============================
def run_experiment(optimizer_name, model_fn, train_loader, test_loader,
                   num_epochs=100, lr=1e-3, momentum=0.9, max_norm=1.0):
    model = model_fn().to(device)
    do_grad_clip = False

    if optimizer_name.lower() == 'fastadam':
        optimizer = FastAdam(model.parameters(), lr=lr, custom=True, weight_decay=1e-4)
    elif optimizer_name.lower() == 'adam':
        optimizer = optim.Adam(model.parameters(), lr=lr, weight_decay=1e-4)
    elif optimizer_name.lower() == 'sgd':
        optimizer = optim.SGD(model.parameters(), lr=lr, momentum=momentum, weight_decay=1e-4)
    elif optimizer_name.lower() == 'adam_gc':
        optimizer = optim.Adam(model.parameters(), lr=lr, weight_decay=1e-4)
        do_grad_clip = True
    else:
        raise ValueError(f"Optimizer '{optimizer_name}' not recognized.")

    history = {
        'train_loss': [],
        'test_loss': [],
        'train_mse': [],
        'test_mse': [],
        'train_ssim': [],
        'test_ssim': []
    }

    for epoch in range(num_epochs):
        model.train()
        running_loss = 0.0
        total = 0
        train_pbar = tqdm(train_loader, desc=f"Epoch {epoch+1}/{num_epochs} Training", leave=False)
        for images, _ in train_pbar:
            images = images.to(device)
            optimizer.zero_grad()
            x_recon, mu, logvar = model(images)
            loss = vae_loss(images, x_recon, mu, logvar, beta=model.beta)
            loss.backward()
            if do_grad_clip:
                clip_grad_norm_(model.parameters(), max_norm)
            optimizer.step()
            running_loss += loss.item()
            total += images.size(0)
            train_pbar.set_postfix(loss=f"{loss.item():.4f}")
        avg_train_loss = running_loss / total
        train_mse = compute_recon_mse(model, train_loader)
        train_ssim = compute_ssim_metric(model, train_loader)

        model.eval()
        running_loss_val = 0.0
        total_val = 0
        val_pbar = tqdm(test_loader, desc=f"Epoch {epoch+1}/{num_epochs} Validation", leave=False)
        with torch.no_grad():
            for images, _ in val_pbar:
                images = images.to(device)
                x_recon, mu, logvar = model(images)
                loss_val = vae_loss(images, x_recon, mu, logvar, beta=model.beta)
                running_loss_val += loss_val.item()
                total_val += images.size(0)
                val_pbar.set_postfix(loss=f"{loss_val.item():.4f}")
        avg_val_loss = running_loss_val / total_val
        test_mse = compute_recon_mse(model, test_loader)
        test_ssim = compute_ssim_metric(model, test_loader)

        history['train_loss'].append(avg_train_loss)
        history['test_loss'].append(avg_val_loss)
        history['train_mse'].append(train_mse)
        history['test_mse'].append(test_mse)
        history['train_ssim'].append(train_ssim)
        history['test_ssim'].append(test_ssim)

        print(f"[{optimizer_name.upper()}] Epoch [{epoch+1}/{num_epochs}] - "
              f"Train Loss: {avg_train_loss:.4f}, Test Loss: {avg_val_loss:.4f}, "
              f"Train MSE: {train_mse:.4f}, Test MSE: {test_mse:.4f}, "
              f"Train SSIM: {train_ssim:.4f}, Test SSIM: {test_ssim:.4f}")

    return model, history

import torch
import torch.nn as nn
import torch.nn.functional as F
import csv  # for logging

# Modified FastAdam optimizer with logging of Eve updates.
class FastAdam:
    def __init__(self, params, lr=0.001, betas=(0.9, 0.999), eps=1e-8,
                 custom=False, weight_decay=1e-4, c=1.0,
                 log_csv="/content/gdrive/My Drive/Colab Notebooks/Eve/eve_updates_log.csv"):
        self.lr = lr
        self.beta1, self.beta2 = betas
        self.eps = eps
        self.custom = custom
        self.weight_decay = weight_decay
        self.t = 0
        self.c = c

        # Open CSV file and write header for aggregated logging.
        if log_csv is not None:
            self.csv_file = open(log_csv, 'w', newline='')
            self.csv_writer = csv.writer(self.csv_file)
            self.csv_writer.writerow(["step", "total_eve_updates", "layer_counts"])
        else:
            self.csv_file = None
            self.csv_writer = None

        # Separate names and tensors.
        self.names = []
        self.params = []
        for i, p in enumerate(params):
            if isinstance(p, tuple):
                self.names.append(p[0])
                self.params.append(p[1])
            else:
                self.names.append("param_" + str(i))
                self.params.append(p)

        # Initialize first and second moment estimates and previous gradients.
        self.m = [torch.zeros_like(p) for p in self.params]
        self.v = [torch.zeros_like(p) for p in self.params]
        self.prev_grad = [torch.zeros_like(p) for p in self.params]

    def zero_grad(self):
        for p in self.params:
            if p.grad is not None:
                p.grad.zero_()

    def step(self):
        self.t += 1

        # Initialize aggregation counters.
        total_eve_update_count = 0
        layer_counts = {}

        for i, (p, m, v, prev_g) in enumerate(zip(self.params, self.m, self.v, self.prev_grad)):
            if p.grad is None:
                continue
            grad = p.grad

            # Apply decoupled weight decay (AdamW style)
            if self.weight_decay != 0:
                p.data.mul_(1 - self.lr * self.weight_decay)

            # Update biased first and second moment estimates.
            m[:] = self.beta1 * m + (1 - self.beta1) * grad
            v[:] = self.beta2 * v + (1 - self.beta2) * grad.square()

            # Compute bias-corrected estimates.
            m_hat = m / (1 - self.beta1 ** self.t)
            v_hat = v / (1 - self.beta2 ** self.t)

            # Compute standard Adam update.
            update = -self.lr * m_hat / (v_hat.sqrt() + self.eps)

            if self.custom:
                # Condition: use Adam update if its magnitude is less than the previous gradient's magnitude.
                condition = update.abs() <= self.c * prev_g.abs()
                # backward_mask is True where the backward update (Eve correction) is applied.
                backward_mask = ~condition
                count = int(backward_mask.sum().item())
                total_eve_update_count += count

                # Extract layer name (using the part before the first dot).
                layer = self.names[i].split('.')[0] if '.' in self.names[i] else self.names[i]
                layer_counts[layer] = layer_counts.get(layer, 0) + count

                # Where condition is False, apply the backward update (using 0.1 * previous gradient).
                update = torch.where(condition, update, -0.1 * prev_g)

            # Update the parameter.
            p.data.add_(update)
            # Store the current gradient for the next iteration.
            prev_g[:] = grad

        # Log aggregated Eve update information for this step.
        if self.csv_writer is not None:
            self.csv_writer.writerow([self.t, total_eve_update_count, str(layer_counts)])

    def close(self):
        if self.csv_file is not None:
            self.csv_file.close()

    def __del__(self):
        self.close()

# ============================
# 6. Run Experiments with Multiple Optimizers for β-VAE
# ============================

save_dir = "/content/gdrive/My Drive/Colab Notebooks/Eve/vae/"

optimizers_list = ['fastadam',
                   #'adam',
                   #'sgd',
                   #'adam_gc'
                   ]
results = {}
final_models = {}

num_epochs = 150
learning_rate = 1e-3
momentum = 0.9
max_norm = 1.0

for opt_name in optimizers_list:
    print(f"\n*** Training with {opt_name.upper()} ***")
    model, history = run_experiment(
        opt_name,
        lambda: create_vae(),
        train_loader,
        test_loader,
        num_epochs=num_epochs,
        lr=learning_rate,
        momentum=momentum,
        max_norm=max_norm
    )
    results[opt_name] = history
    final_models[opt_name] = model
    # Salva il modello (solo lo state_dict per maggiore compatibilità)
    model_save_path = os.path.join(save_dir, f"{opt_name}_{dataset_name}_{str(beta)}_model_state_dict_-.pth")
    torch.save(model.state_dict(), model_save_path)
    print(f"Modello {opt_name} salvato in: {model_save_path}")