From a5d5f30fbd9c6c7c78834072401932c84bddaf14 Mon Sep 17 00:00:00 2001 From: gdamms Date: Thu, 5 Feb 2026 15:33:20 +0100 Subject: trying to improve whole project --- .gitignore | 177 +----------- README.md | 81 +++++- autoencoder.py | 154 ----------- main.py | 695 +++++++++-------------------------------------- models/__init__.py | 6 + models/autoencoder.py | 103 +++++++ models/unet.py | 140 ++++++++++ plots.py | 174 ------------ requirements.txt | 11 +- src/__init__.py | 50 ++++ src/config.py | 34 +++ src/dataloader.py | 210 ++++++++++++++ src/diffusion.py | 154 +++++++++++ src/metrics.py | 98 +++++++ src/sample.py | 248 +++++++++++++++++ src/train_autoencoder.py | 167 ++++++++++++ src/train_diffusion.py | 197 ++++++++++++++ src/utils.py | 110 ++++++++ utils.py | 108 -------- 19 files changed, 1739 insertions(+), 1178 deletions(-) delete mode 100644 autoencoder.py create mode 100644 models/__init__.py create mode 100644 models/autoencoder.py create mode 100644 models/unet.py delete mode 100644 plots.py create mode 100644 src/__init__.py create mode 100644 src/config.py create mode 100644 src/dataloader.py create mode 100644 src/diffusion.py create mode 100644 src/metrics.py create mode 100644 src/sample.py create mode 100644 src/train_autoencoder.py create mode 100644 src/train_diffusion.py create mode 100644 src/utils.py delete mode 100644 utils.py diff --git a/.gitignore b/.gitignore index 0c22bca..c1eb923 100644 --- a/.gitignore +++ b/.gitignore @@ -1,172 +1,17 @@ -# Tensorboard -runs/ +# Python virtual environment +.venv/ -# Torch datas. -data/ +# Plots +plots/ -# Model weights. -*.pth +# Checkpoints +checkpoints/ -# Tmp files. -*.tmp.* +# MLflow artifacts +mlflow.db -# Byte-compiled / optimized / DLL files +# Python bytecode __pycache__/ -*.py[cod] -*$py.class - -# C extensions -*.so - -# Distribution / packaging -.Python -build/ -develop-eggs/ -dist/ -downloads/ -eggs/ -.eggs/ -lib/ -lib64/ -parts/ -sdist/ -var/ -wheels/ -share/python-wheels/ -*.egg-info/ -.installed.cfg -*.egg -MANIFEST - -# PyInstaller -# Usually these files are written by a python script from a template -# before PyInstaller builds the exe, so as to inject date/other infos into it. -*.manifest -*.spec - -# Installer logs -pip-log.txt -pip-delete-this-directory.txt - -# Unit test / coverage reports -htmlcov/ -.tox/ -.nox/ -.coverage -.coverage.* -.cache -nosetests.xml -coverage.xml -*.cover -*.py,cover -.hypothesis/ -.pytest_cache/ -cover/ - -# Translations -*.mo -*.pot - -# Django stuff: -*.log -local_settings.py -db.sqlite3 -db.sqlite3-journal - -# Flask stuff: -instance/ -.webassets-cache - -# Scrapy stuff: -.scrapy - -# Sphinx documentation -docs/_build/ - -# PyBuilder -.pybuilder/ -target/ - -# Jupyter Notebook -.ipynb_checkpoints - -# IPython -profile_default/ -ipython_config.py - -# pyenv -# For a library or package, you might want to ignore these files since the code is -# intended to run in multiple environments; otherwise, check them in: -# .python-version -# pipenv -# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. -# However, in case of collaboration, if having platform-specific dependencies or dependencies -# having no cross-platform support, pipenv may install dependencies that don't work, or not -# install all needed dependencies. -#Pipfile.lock - -# poetry -# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. -# This is especially recommended for binary packages to ensure reproducibility, and is more -# commonly ignored for libraries. -# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control -#poetry.lock - -# pdm -# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. -#pdm.lock -# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it -# in version control. -# https://pdm.fming.dev/#use-with-ide -.pdm.toml - -# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm -__pypackages__/ - -# Celery stuff -celerybeat-schedule -celerybeat.pid - -# SageMath parsed files -*.sage.py - -# Environments -.env -.venv -env/ -venv/ -ENV/ -env.bak/ -venv.bak/ - -# Spyder project settings -.spyderproject -.spyproject - -# Rope project settings -.ropeproject - -# mkdocs documentation -/site - -# mypy -.mypy_cache/ -.dmypy.json -dmypy.json - -# Pyre type checker -.pyre/ - -# pytype static type analyzer -.pytype/ - -# Cython debug symbols -cython_debug/ - -# PyCharm -# JetBrains specific template is maintained in a separate JetBrains.gitignore that can -# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore -# and can be added to the global gitignore or merged into this file. For a more nuclear -# option (not recommended) you can uncomment the following to ignore the entire idea folder. -#.idea/ +# Data +data/ diff --git a/README.md b/README.md index 93cdc90..805baae 100644 --- a/README.md +++ b/README.md @@ -1,2 +1,79 @@ -# diffusion-mnist -Diffusion inspired mnist like image generation. +# MNIST Diffusion Model + +A diffusion-based generative model for MNIST digits implemented in PyTorch. + +## Project Structure + +``` +diffusion-mnist/ +├── main.py # Main entry point with CLI +├── models/ # Neural network architectures +│ ├── __init__.py +│ ├── unet.py # UNet for diffusion model +│ └── autoencoder.py # Autoencoder for latent diffusion +├── src/ # Source code modules +│ ├── __init__.py +│ ├── config.py # Configuration and hyperparameters +│ ├── diffusion.py # Diffusion process utilities +│ ├── dataloader.py # Dataset and dataloader classes +│ ├── utils.py # Helper functions and metrics +│ ├── train_diffusion.py # Diffusion training script +│ ├── train_autoencoder.py # Autoencoder training script +│ └── sample.py # Sampling and visualization +├── checkpoints/ # Model checkpoints +├── plots/ # Generated visualizations +├── data/ # Dataset directory +└── runs/ # TensorBoard logs +``` + +## Installation + +```bash +pip install -r requirements.txt +``` + +## Usage + +### Train Diffusion Model +```bash +python main.py train --epochs 10 --lr 2e-4 --batch-size 64 +``` + +### Train with Self-Attention +```bash +python main.py train --epochs 10 --attention +``` + +### Train Autoencoder (for latent diffusion) +```bash +python main.py train-ae --epochs 10 +``` + +### Generate Samples +```bash +python main.py sample --checkpoint checkpoints/diffusion_latest.pt +``` + +### Visualize Diffusion Process +```bash +python main.py visualize --all +``` + +## Configuration + +All hyperparameters can be found in `src/config.py`: +- `DIFFU_STEPS`: Number of diffusion steps (default: 1000) +- `EPOCHS`: Training epochs (default: 10) +- `BATCH_SIZE`: Batch size (default: 64) +- `LEARNING_RATE`: Learning rate (default: 2e-4) + +## Model Architecture + +The diffusion model uses a UNet architecture with: +- Timestep embedding +- Label conditioning (for class-conditional generation) +- Optional self-attention layers + +## License + +See [LICENSE](LICENSE) for details. diff --git a/autoencoder.py b/autoencoder.py deleted file mode 100644 index 23c47f9..0000000 --- a/autoencoder.py +++ /dev/null @@ -1,154 +0,0 @@ -import torch -from torch.utils.data import DataLoader, Dataset - -from torchvision import datasets, transforms - -import matplotlib.pyplot as plt - -from trainer import train - - -class PrintLayer(torch.nn.Module): - def forward(self, x): - print(x.shape) - print(x.min()) - print(x.max()) - return x - - -class Autoencoder(torch.nn.Module): - def __init__(self, input_dim, latent_dim): - super(Autoencoder, self).__init__() - self.input_dim = input_dim - self.latent_dim = latent_dim - self.encoder = torch.nn.Sequential( - torch.nn.Conv2d(input_dim, 16, kernel_size=3, padding=1), - torch.nn.ReLU(), - torch.nn.Conv2d(16, 16, kernel_size=3, padding=1), - torch.nn.ReLU(), - torch.nn.MaxPool2d(kernel_size=2), - torch.nn.Conv2d(16, 32, kernel_size=3, padding=1), - torch.nn.ReLU(), - torch.nn.Conv2d(32, 32, kernel_size=3, padding=1), - torch.nn.ReLU(), - torch.nn.MaxPool2d(kernel_size=2), - torch.nn.Conv2d(32, 64, kernel_size=3, padding=1), - torch.nn.ReLU(), - torch.nn.Conv2d(64, 64, kernel_size=3, padding=1), - torch.nn.ReLU(), - torch.nn.Conv2d(64, latent_dim, kernel_size=3, padding=1), - torch.nn.ReLU(), - # 1x7x7 to 1x8x8 - torch.nn.Conv2d(latent_dim, latent_dim, kernel_size=2, padding=1), - torch.nn.Sigmoid(), - ) - self.decoder = torch.nn.Sequential( - # 1x8x8 to 1x7x7 - torch.nn.Conv2d(latent_dim, latent_dim, kernel_size=2, padding=0), - torch.nn.ReLU(), - # main decoder - torch.nn.Conv2d(latent_dim, 64, kernel_size=3, padding=1), - torch.nn.ReLU(), - torch.nn.Conv2d(64, 64, kernel_size=3, padding=1), - torch.nn.ReLU(), - torch.nn.ConvTranspose2d(64, 32, kernel_size=2, stride=2), - torch.nn.ReLU(), - torch.nn.Conv2d(32, 32, kernel_size=3, padding=1), - torch.nn.ReLU(), - torch.nn.ConvTranspose2d(32, 16, kernel_size=2, stride=2), - torch.nn.ReLU(), - torch.nn.Conv2d(16, 16, kernel_size=3, padding=1), - torch.nn.ReLU(), - torch.nn.Conv2d(16, input_dim, kernel_size=3, padding=1), - torch.nn.Sigmoid(), - ) - - def forward(self, x): - x = self.encoder(x) - x = self.decoder(x) - return x - - def encode(self, x): - return self.encoder(x) - - def decode(self, x): - return self.decoder(x) - - -class AutoencoderDataset(Dataset): - def __init__(self, dataset, device='cpu'): - self.dataset = dataset - self.device = device - self.dummy_param = torch.nn.Parameter(torch.empty(0)) - - def __len__(self): - return len(self.dataset) - - def __getitem__(self, idx): - data = self.dataset[idx][0].to(self.device) - return data, data - -def main(): - torch.multiprocessing.set_start_method("spawn") - device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') - - # Load dataset - mnist = datasets.MNIST( - root='data', - train=True, - download=True, - transform=transforms.ToTensor(), - ) - dataset = AutoencoderDataset(mnist, device=device) - dataloader = DataLoader(dataset, batch_size=64, shuffle=True, - num_workers=4, persistent_workers=True) - - # Initialize model - model = Autoencoder(input_dim=1, latent_dim=1) - model.load_state_dict(torch.load('mnist_autoencoder.pth')) - model.to(device) - - # Train model - lr = 1e-3 - epochs = 1 - optimizer = torch.optim.Adam(model.parameters(), lr=lr) - criterion = torch.nn.functional.binary_cross_entropy - train(model, dataloader, epochs, optimizer, criterion) - - # Save model - torch.save(model.state_dict(), 'autoencoder.pth') - - # Visualize results - n = 10 - with torch.no_grad(): - plt.figure(figsize=(2*n, 6)) - for i, j in enumerate(torch.randint(0, len(dataset), (n,))): - x, _ = dataset[j] - x = x.unsqueeze(0) - x_latent = model.encode(x) - x_hat = model.decode(x_latent) - plt.subplot(3, n+1, i + 2) - plt.imshow(x.cpu().squeeze().numpy()) - plt.axis('off') - plt.subplot(3, n+1, i + n + 3) - plt.imshow(x_latent.cpu().squeeze().numpy()) - plt.axis('off') - plt.subplot(3, n+1, i + 2*n + 4) - plt.imshow(x_hat.cpu().squeeze().numpy()) - plt.axis('off') - plt.subplot(3, n+1, 1) - plt.text(0.5, 0.5, 'Original', horizontalalignment='center', fontsize=12) - plt.axis('off') - plt.subplot(3, n+1, n + 2) - plt.text(0.5, 0.5, 'Latent', horizontalalignment='center', fontsize=12) - plt.axis('off') - plt.subplot(3, n+1, 2*n + 3) - plt.text(0.5, 0.5, 'Reconstructed', horizontalalignment='center', fontsize=12) - plt.axis('off') - plt.suptitle('Autoencoder') - plt.tight_layout() - plt.savefig('autoencoder.tmp.png') - - -if __name__ == '__main__': - main() diff --git a/main.py b/main.py index 7f8444a..46c8a50 100644 --- a/main.py +++ b/main.py @@ -1,580 +1,137 @@ -import torch -import torch.nn as nn -import torch.nn.functional as F -import torch.nn.attention as attention -from torch.utils.data import DataLoader, Dataset -from rich.progress import track -import io -from PIL import Image - -from torchvision import datasets, transforms - -from trainer import train -from trainer.trainer import Trainer - -import matplotlib.pyplot as plt - -import numpy as np -import os -import cv2 - -from autoencoder import Autoencoder -from utils import * - - -class SelfAttention(nn.Module): - def __init__(self, nb_channels, nb_heads): - super().__init__() - self.attention = nn.MultiheadAttention(nb_channels, nb_heads) - - def forward(self, x): - _, c, w, h = x.shape - x = x.view(-1, c, w*h) - x = x.permute(2, 0, 1) - x, _ = self.attention(x, x, x) - x = x.permute(1, 2, 0) - x = x.view(-1, c, w, h) - return x - - -class UNetAddAttUEDF(nn.Module): - def __init__(self): - super().__init__() - - ## Inputs: - # xt: image at step t (NB_CHANNEL*IMG_SIZE*IMG_SIZE) - # t: step number (1) - # vec: one-hot vector of the label (NB_LABEL) - - ## Encoder for t - self.encodet = nn.Linear(1, IMG_SIZE*IMG_SIZE) - - ## Encoder for vec - self.encodevec = nn.Linear(NB_LABEL, IMG_SIZE*IMG_SIZE) - - ## UNet (2 more channels input because we concatenate xt with t and vec) - self.conv1 = nn.Conv2d(NB_CHANNEL+2, 64, 3, padding=1) - self.conv2 = nn.Conv2d(64, 64, 3, padding=1) - self.maxpool1 = nn.MaxPool2d(2, 2) - self.conv3 = nn.Conv2d(64, 128, 3, padding=1) - self.conv4 = nn.Conv2d(128, 128, 3, padding=1) - self.maxpool2 = nn.MaxPool2d(2, 2) - self.att1 = SelfAttention(128, 8) - self.conv5 = nn.Conv2d(128, 256, 3, padding=1) - self.att2 = SelfAttention(256, 8) - self.conv6 = nn.Conv2d(256, 256, 3, padding=1) - self.att3 = SelfAttention(256, 8) - self.upconv1 = nn.ConvTranspose2d(256, 128, 2, stride=2) - self.conv7 = nn.Conv2d(256, 128, 3, padding=1) - self.conv8 = nn.Conv2d(128, 128, 3, padding=1) - self.upconv2 = nn.ConvTranspose2d(128, 64, 2, stride=2) - self.conv9 = nn.Conv2d(128, 64, 3, padding=1) - self.conv10 = nn.Conv2d(64, 64, 3, padding=1) - self.conv11 = nn.Conv2d(64, NB_CHANNEL, 3, padding=1) - - def forward(self, xt, t, vec): - # Encode t and vec - t = F.relu(self.encodet(t / DIFFU_STEPS)) - t = t.view(-1, 1, IMG_SIZE, IMG_SIZE) - vec = F.relu(self.encodevec(vec)) - vec = vec.view(-1, 1, IMG_SIZE, IMG_SIZE) - - # Concat all 3 - x = torch.cat((xt, t, vec), dim=1) - - # UNet - x1 = F.relu(self.conv1(x)) - x1 = F.relu(self.conv2(x1)) - x2 = self.maxpool1(x1) - x2 = F.relu(self.conv3(x2)) - x2 = F.relu(self.conv4(x2)) - x3 = self.maxpool2(x2) - x3 = self.att1(x3) - x3 = F.relu(self.conv5(x3)) - x3 = self.att2(x3) - x3 = F.relu(self.conv6(x3)) - x4 = self.upconv1(x3) - x4 = torch.cat((x4, x2), dim=1) - x4 = F.relu(self.conv7(x4)) - x4 = F.relu(self.conv8(x4)) - x5 = self.upconv2(x4) - x5 = torch.cat((x5, x1), dim=1) - x5 = F.relu(self.conv9(x5)) - x5 = F.relu(self.conv10(x5)) - x5 = self.conv11(x5) - - return x5 - - -class UNetMNISTSig(nn.Module): - def __init__(self): - super().__init__() - - ## Inputs: - # xt: image at step t (NB_CHANNEL*IMG_SIZE*IMG_SIZE) - # t: step number (1) - # vec: one-hot vector of the label (NB_LABEL) - - ## Encoder for t - self.encodet = nn.Linear(1, IMG_SIZE*IMG_SIZE) - - ## Encoder for vec - self.encodevec = nn.Linear(NB_LABEL, IMG_SIZE*IMG_SIZE) - - ## UNet (2 more channels input because we concatenate xt with t and vec) - self.conv1 = nn.Conv2d(NB_CHANNEL+2, 64, 3, padding=1) - self.conv2 = nn.Conv2d(64, 64, 3, padding=1) - self.maxpool1 = nn.MaxPool2d(2, 2) - self.conv3 = nn.Conv2d(64, 128, 3, padding=1) - self.conv4 = nn.Conv2d(128, 128, 3, padding=1) - self.maxpool2 = nn.MaxPool2d(2, 2) - self.conv5 = nn.Conv2d(128, 256, 3, padding=1) - self.conv6 = nn.Conv2d(256, 256, 3, padding=1) - self.upconv1 = nn.ConvTranspose2d(256, 128, 2, stride=2) - self.conv7 = nn.Conv2d(256, 128, 3, padding=1) - self.conv8 = nn.Conv2d(128, 128, 3, padding=1) - self.upconv2 = nn.ConvTranspose2d(128, 64, 2, stride=2) - self.conv9 = nn.Conv2d(128, 64, 3, padding=1) - self.conv10 = nn.Conv2d(64, 64, 3, padding=1) - self.conv11 = nn.Conv2d(64, NB_CHANNEL, 3, padding=1) - - def forward(self, xt, t, vec): - # Encode t and vec - t = F.relu(self.encodet(t / DIFFU_STEPS)) - t = t.view(-1, 1, IMG_SIZE, IMG_SIZE) - vec = F.relu(self.encodevec(vec)) - vec = vec.view(-1, 1, IMG_SIZE, IMG_SIZE) - - # Concat all 3 - x = torch.cat((xt, t, vec), dim=1) - - # UNet - x1 = F.relu(self.conv1(x)) - x1 = F.relu(self.conv2(x1)) - x2 = self.maxpool1(x1) - x2 = F.relu(self.conv3(x2)) - x2 = F.relu(self.conv4(x2)) - x3 = self.maxpool2(x2) - x3 = F.relu(self.conv5(x3)) - x3 = F.relu(self.conv6(x3)) - x4 = self.upconv1(x3) - x4 = torch.cat((x4, x2), dim=1) - x4 = F.relu(self.conv7(x4)) - x4 = F.relu(self.conv8(x4)) - x5 = self.upconv2(x4) - x5 = torch.cat((x5, x1), dim=1) - x5 = F.relu(self.conv9(x5)) - x5 = F.relu(self.conv10(x5)) - x5 = self.conv11(x5) - - return x5 - -class FolderDataset(Dataset): - def __init__(self, path, size=(32, 32)): - super().__init__() - self.path = path - self.size = size - self.files = os.listdir(self.path) +""" +MNIST Diffusion Model - def __getitem__(self, index): - img = cv2.imread(os.path.join(self.path, self.files[index])) - img = cv2.resize(img, self.size) - img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) - img = np.transpose(img, (2, 0, 1)) / 255 - return torch.tensor(img, dtype=torch.float32), 0 +A diffusion-based generative model for MNIST digits. - def __len__(self): - return len(self.files) +Usage: + Train diffusion model: + python main.py train --epochs 10 + Train autoencoder: + python main.py train-ae --epochs 10 -def q_xt_xt_1(xt_1, t): - t_ind = t.to(dtype=torch.long) if isinstance(t, torch.Tensor) else t + Generate samples: + python main.py sample --checkpoint checkpoints/diffusion_latest.pt - alpha = ALPHA[t_ind] - mean = torch.sqrt(alpha) * xt_1 - std = torch.sqrt(1 - alpha) + Visualize diffusion process: + python main.py visualize --all +""" - eps = torch.randn(xt_1.shape, device=DEVICE) - xt = mean + std * eps - - return xt, eps - - -def q_xt_x0(x0, t): - t_ind = t.to(dtype=torch.long) if isinstance(t, torch.Tensor) else t - - reshaped = len(x0.shape) == 3 - if reshaped: - c, w, h = x0.shape - b = 1 - x0 = x0.view(b, c, w, h) - else: - b, c, w, h = x0.shape - t_ind = t_ind.view(b, 1, 1, 1) - t_ind = t_ind.expand(b, c, w, h) - - alpha_bar = ALPHA_BAR[t_ind] - mean = torch.sqrt(alpha_bar) * x0 - std = torch.sqrt(1 - alpha_bar) - - eps = torch.randn(x0.shape, device=DEVICE) - xt = mean + std * eps - - if reshaped: - xt = xt.view(c, w, h) - - return xt, eps - - -def p_xt_1_xt(model, xt, t, vec): - t_ind = t.to(dtype=torch.long) if isinstance(t, torch.Tensor) else t - - alpha_bar_t = ALPHA_BAR[t_ind].view(-1, 1, 1, 1) - alpha_bar_t_1 = ALPHA_BAR[t_ind-1].view(-1, 1, 1, 1) - alpha_t = ALPHA[t_ind].view(-1, 1, 1, 1) - beta_t = BETA[t_ind].view(-1, 1, 1, 1) - - beta_tilde = (1 - alpha_bar_t_1) / (1 - alpha_bar_t) * beta_t - - epsilon_theta = model(xt, t, vec) - - sigma_theta = torch.sqrt(beta_tilde) - mu_theta = (xt - (1 - alpha_t) / torch.sqrt(1 - alpha_bar_t) * epsilon_theta) / torch.sqrt(alpha_t) - - mask_t0 = (t > 1).to(dtype=torch.float32).view(-1, 1, 1, 1) - noise = torch.randn(xt.shape, device=DEVICE) * mask_t0 - - return mu_theta + sigma_theta * noise - - -def p_xt_1_xt_sig(model, xt_1, t, vec): - x0 = model(xt_1, t, vec) - xt, _ = q_xt_x0(x0, t) - return xt - - -class DiffusionDataset(Dataset): - def __init__(self, dataset, autoencoder=None): - super().__init__() - self.dataset = dataset - self.autoencoder = autoencoder - - def __getitem__(self, index): - # Get the image and the label. - img, label = self.dataset[index] - img = img.to(DEVICE) - - # Encode the image. - if self.autoencoder is not None: - img = self.autoencoder.encode(img.unsqueeze(0)).squeeze(0) - - # Normalize the image. - img = img * 2 - 1 - - # Add noise to the image. - t = torch.randint(1, DIFFU_STEPS, (1,), device=DEVICE) - xt, eps = q_xt_x0(img, t) +import argparse +import torch - # Convert the label to a one-hot vector. - vec = torch.nn.functional.one_hot( - torch.tensor(min(label, NB_LABEL-1)), - num_classes=NB_LABEL, - ) - return ( - # x_true - xt.clone().detach().to(dtype=torch.float32, device=DEVICE), - t.clone().detach().to(dtype=torch.float32, device=DEVICE), - vec.clone().detach().to(dtype=torch.float32, device=DEVICE), - # y_true - eps, +def main(): + parser = argparse.ArgumentParser( + description="MNIST Diffusion Model", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=__doc__ + ) + + subparsers = parser.add_subparsers(dest="command", help="Available commands") + + # Train diffusion model + train_parser = subparsers.add_parser("train", help="Train diffusion model") + train_parser.add_argument("--epochs", type=int, default=10, help="Number of epochs") + train_parser.add_argument("--lr", type=float, default=2e-4, help="Learning rate") + train_parser.add_argument("--batch-size", type=int, default=64, help="Batch size") + train_parser.add_argument("--attention", action="store_true", help="Use self-attention") + train_parser.add_argument("--checkpoint", type=str, default=None, help="Resume from checkpoint") + train_parser.add_argument("--name", type=str, default=None, help="Run name") + + # Train autoencoder + ae_parser = subparsers.add_parser("train-ae", help="Train autoencoder") + ae_parser.add_argument("--epochs", type=int, default=10, help="Number of epochs") + ae_parser.add_argument("--lr", type=float, default=1e-3, help="Learning rate") + ae_parser.add_argument("--batch-size", type=int, default=64, help="Batch size") + ae_parser.add_argument("--latent-channels", type=int, default=1, help="Latent channels") + ae_parser.add_argument("--checkpoint", type=str, default=None, help="Resume from checkpoint") + + # Sample from model + sample_parser = subparsers.add_parser("sample", help="Generate samples") + sample_parser.add_argument("--checkpoint", type=str, default="checkpoints/diffusion_latest.pt", + help="Path to model checkpoint") + sample_parser.add_argument("--n-samples", type=int, default=10, help="Samples per class") + sample_parser.add_argument("--attention", action="store_true", help="Use attention in model") + + # Visualize diffusion + viz_parser = subparsers.add_parser("visualize", help="Visualize diffusion process") + viz_parser.add_argument("--checkpoint", type=str, default="checkpoints/diffusion_latest.pt", + help="Path to model checkpoint") + viz_parser.add_argument("--attention", action="store_true", help="Use attention in model") + viz_parser.add_argument("--forward", action="store_true", help="Visualize forward diffusion") + viz_parser.add_argument("--backward", action="store_true", help="Visualize backward diffusion") + viz_parser.add_argument("--all", action="store_true", help="Run all visualizations") + + args = parser.parse_args() + + if args.command is None: + parser.print_help() + return + + # Set multiprocessing start method + torch.multiprocessing.set_start_method("spawn", force=True) + + if args.command == "train": + from src.train_diffusion import train_diffusion + train_diffusion( + epochs=args.epochs, + learning_rate=args.lr, + batch_size=args.batch_size, + use_attention=args.attention, + checkpoint_path=args.checkpoint, + run_name=args.name, ) - - def __len__(self): - return len(self.dataset) - - -class DiffusionDatasetSig(Dataset): - def __init__(self, dataset, autoencoder=None): - super().__init__() - self.dataset = dataset - self.autoencoder = autoencoder - - def __getitem__(self, index): - # Get the image and the label. - img, label = self.dataset[index] - img = img.to(DEVICE) - - # Encode the image. - if self.autoencoder is not None: - img = self.autoencoder.encode(img.unsqueeze(0)).squeeze(0) - - # Normalize the image. - img = img * 2 - 1 - - # Add noise to the image. - t = torch.randint(1, DIFFU_STEPS, (1,), device=DEVICE) - xt, eps = q_xt_x0(img, t) - - # Convert the label to a one-hot vector. - vec = torch.nn.functional.one_hot( - torch.tensor(min(label, NB_LABEL-1)), - num_classes=NB_LABEL, + + elif args.command == "train-ae": + from src.train_autoencoder import train_autoencoder + train_autoencoder( + epochs=args.epochs, + learning_rate=args.lr, + batch_size=args.batch_size, + latent_channels=args.latent_channels, + checkpoint_path=args.checkpoint, ) - - return ( - # x_true - xt.clone().detach().to(dtype=torch.float32, device=DEVICE), - t.clone().detach().to(dtype=torch.float32, device=DEVICE), - vec.clone().detach().to(dtype=torch.float32, device=DEVICE), - # y_true - img.clone().detach().to(dtype=torch.float32, device=DEVICE), + + elif args.command == "sample": + import os + from src.config import DEVICE + from src.sample import generate_grid + from src.utils import load_checkpoint + from models import UNetMNIST + + model = UNetMNIST(use_attention=args.attention).to(DEVICE) + if os.path.exists(args.checkpoint): + model = load_checkpoint(model, os.path.basename(args.checkpoint)) + else: + print(f"Warning: Checkpoint {args.checkpoint} not found.") + + generate_grid(model, n_per_class=args.n_samples) + + elif args.command == "visualize": + import os + from src.config import DEVICE + from src.sample import ( + visualize_forward_diffusion, + visualize_backward_diffusion, + generate_grid, ) - - def __len__(self): - return len(self.dataset) - - -def loss(y_pred, y_true): - return nn.MSELoss()(y_pred, y_true) - - -def forward_diffusion(x0): - x = x0.clone() - xs = [x] - for t in range(1, DIFFU_STEPS+1): - x, _ = q_xt_xt_1(x, t) - xs.append(x) - return xs - - -def tensor_to_image(tensor): - img = tensor.clone().detach().cpu().numpy().transpose(1, 2, 0) - img -= img.min() - img /= img.max() - return img - - -def tensor_to_images(tensor): - img = tensor.clone().detach().cpu().numpy().transpose(0, 2, 3, 1) - img -= np.min(img, axis=(1, 2, 3), keepdims=True) - img /= np.max(img, axis=(1, 2, 3), keepdims=True) - return img - - -def figure_to_image(figure): - buf = io.BytesIO() - figure.savefig(buf, format='png') - buf.seek(0) - image = np.array(Image.open(buf)) - return image - - -DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") - -DIFFU_STEPS = 1000 -BETA = torch.linspace(1e-4, 2e-2, DIFFU_STEPS, device=DEVICE) -BETA = torch.cat((torch.tensor([0.], device=DEVICE), BETA)) -ALPHA = 1 - BETA -ALPHA_BAR = torch.cumprod(ALPHA, dim=0) - -dataset = datasets.MNIST( - root="./data", - train=True, - download=True, - transform=transforms.ToTensor(), -) -# dataset = datasets.LFWPeople( -# root="./data", -# download=True, -# transform=transforms.Compose([ -# transforms.Resize((64, 64)), -# transforms.ToTensor(), -# ]), -# ) -# dataset = FolderDataset('data/lfwcrop_color/faces') -# dataset = FolderDataset('data/edface') - -autoencoder = None -# autoencoder = Autoencoder(1, 1).to(DEVICE) -# autoencoder.load_state_dict(torch.load('autoencoder.pth')) -# autoencoder.eval() - -img = dataset[0][0].to(DEVICE) -if autoencoder is not None: - img = autoencoder.encode(img.unsqueeze(0)).squeeze(0) -NB_CHANNEL, IMG_SIZE, _ = img.shape -NB_LABEL = 10 - -EPOCHS = 10 -LEARNING_RATE = 2e-4 - - -def epoch_callback(trainer: Trainer): - epoch_i = trainer.epoch_i - - if epoch_i % 1 == 0 or epoch_i == trainer.epoch_end: - save_path = f'runs/{trainer.run_name}/checkpoints/{epoch_i:04}e.pt' - torch.save(trainer.model, save_path) - save_path = f'runs/{trainer.run_name}/checkpoints/last.pt' - torch.save(trainer.model, save_path) - - print("Calculating metrics...") - with torch.no_grad(): - batch_size = 64 - n_batches = 16 - n_samples = batch_size * n_batches - - fakes = np.zeros((0, NB_CHANNEL, IMG_SIZE, IMG_SIZE)) - for _ in range(n_batches): - x = torch.randn(batch_size, NB_CHANNEL, IMG_SIZE, IMG_SIZE).to(DEVICE) - vec = torch.randint(0, NB_LABEL, (batch_size,)).to(DEVICE) - vec = torch.nn.functional.one_hot(vec, num_classes=NB_LABEL).to(device=DEVICE, dtype=torch.float32) - - for t in range(DIFFU_STEPS, 0, -1): - t_tensor = torch.tensor([[t]] * batch_size, device=DEVICE, dtype=torch.float32) - x = p_xt_1_xt_sig(model, x, t_tensor, vec) - - x = x.cpu().numpy() - x -= x.min(axis=(1, 2, 3), keepdims=True) - x /= x.max(axis=(1, 2, 3), keepdims=True) - fakes = np.concatenate((fakes, x)) - - reals = torch.stack([dataset[i][0] for i in range(n_samples)]).cpu().numpy() - reals = reals * 2 - 1 - - trainer.writer.add_scalar('FID/Validation', fid(reals, fakes), epoch_i) - trainer.writer.add_scalar('KL/Validation', kl(reals, fakes), epoch_i) - trainer.writer.add_scalar('RKL/Validation', kl(fakes, reals), epoch_i) - trainer.writer.add_scalar('JSD/Validation', jsd(reals, fakes), epoch_i) - - - fig = plt.figure(figsize=(32, 16)) - for i in range(32): - plt.subplot(4, 8, i + 1) - plt.imshow(fakes[i].transpose(1, 2, 0)) - plt.axis("off") - plt.tight_layout() - trainer.writer.add_image('Fakes/Validation', figure_to_image(fig), epoch_i, dataformats='HWC') - plt.close(fig) - - -if __name__ == '__main__': - - ############ - # Training # - ############ - - torch.multiprocessing.set_start_method("spawn") - - # Load the model. - # model = UNetMNISTSig().to(DEVICE) - model = torch.load('runs/20240719-154613_UNetMNISTSig/checkpoints/last.pt').to(DEVICE) - - # Define the optimizer. - optimizer = torch.optim.Adam(model.parameters(), lr=LEARNING_RATE) - - # Define the training dataset. - train_dataset = DiffusionDatasetSig(dataset, autoencoder) - train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True, - num_workers=4, persistent_workers=True) - criterion = loss - epochs = EPOCHS - - # Train the model. - train(model, train_loader, epochs, optimizer, criterion, epoch_callbacks=[epoch_callback], save_chekpoint=False) - - ############## - # Evaluation # - ############## - - with torch.no_grad(): - # Forward diffusion - img, label = dataset[np.random.randint(0, len(dataset))] - img = img.to(DEVICE) - if autoencoder is not None: - img = autoencoder.encode(img.unsqueeze(0)).squeeze(0) - img = img * 2 - 1 - - nb_plots = 10 - plots_id = [i for i in np.linspace(1, DIFFU_STEPS, nb_plots, dtype=int)] - - xs = forward_diffusion(img) - - plt.figure(figsize=(nb_plots, 2.5)) - for plot_i, t in enumerate(plots_id): - x = xs[t] - plt.subplot(2, nb_plots + 1, plot_i + 2) - plt.title(f"t={t}") - plt.imshow(tensor_to_image(x), interpolation='none') - plt.axis("off") - - for plot_i, t in enumerate(plots_id): - x, _ = q_xt_x0(img, t) - plt.subplot(2, nb_plots + 1, nb_plots + plot_i + 3) - plt.imshow(tensor_to_image(x), interpolation='none') - plt.axis("off") - - plt.subplot(2, nb_plots + 1, 1) - plt.text(0, 0.5, "Implicit", fontsize=12) - plt.axis("off") - plt.subplot(2, nb_plots + 1, nb_plots + 2) - plt.text(0, 0.5, "Explicit", fontsize=12) - plt.axis("off") - - plt.suptitle("Forward diffusion") - plt.tight_layout() - plt.savefig("plots/forward_diffusion.tmp.png") - - - # Backward diffusion - t_plots = np.linspace(1, DIFFU_STEPS, nb_plots, dtype=int) - - n_classes = 10 - - x = torch.randn(n_classes, NB_CHANNEL, IMG_SIZE, IMG_SIZE, device=DEVICE) - vec = torch.tensor([[min(i, NB_LABEL-1)] for i in range(n_classes)], dtype=torch.int64) - vec = torch.nn.functional.one_hot(vec, num_classes=NB_LABEL).to(device=DEVICE, dtype=torch.float32) - - plt.figure(figsize=(nb_plots, n_classes)) - plt.suptitle("Backward diffusion") - for t in track(range(DIFFU_STEPS, 0, -1), description='Diffusing...'): - t_tensor = torch.tensor([[t]] * n_classes, device=DEVICE, dtype=torch.float32) - x = p_xt_1_xt_sig(model, x, t_tensor, vec) - if t in t_plots: - t_plot_i = nb_plots - t_plots.tolist().index(t) - 1 - for class_i in range(n_classes): - plt.subplot(n_classes, nb_plots, t_plot_i + nb_plots * class_i + 1) - if class_i == 0: - plt.title(f"t={t}") - plt.imshow(tensor_to_image(x[class_i])) - plt.axis("off") - plt.tight_layout() - plt.savefig("plots/backward_diffusion.tmp.png") - - - # Benchmark - x = torch.randn(nb_plots * n_classes, NB_CHANNEL, IMG_SIZE, IMG_SIZE).to(DEVICE) - vec = sum([[[min(i, NB_LABEL-1)]] * nb_plots for i in range(n_classes)], []) - vec = torch.tensor(vec, device=DEVICE, dtype=torch.int64) - vec = torch.nn.functional.one_hot(vec, num_classes=NB_LABEL).to(device=DEVICE, dtype=torch.float32) - - for ti in track(range(DIFFU_STEPS, 0, -1), description='Benchmarking...'): - t = torch.tensor([[ti]] * n_classes * nb_plots, device=DEVICE, dtype=torch.float32) - x = p_xt_1_xt_sig(model, x, t, vec) - - x = x * 0.5 + 0.5 - x = x.clamp(0, 1) - - plt.figure(figsize=(nb_plots, n_classes)) - for i in range(nb_plots): - for j in range(n_classes): - id = i * n_classes + j - img = x[id] - if autoencoder is not None: - img = autoencoder.decode(img.unsqueeze(0)).squeeze(0) - plt.subplot(n_classes, nb_plots, id + 1) - plt.imshow(tensor_to_image(img)) - plt.axis("off") - plt.tight_layout() - plt.savefig("plots/benchmark.tmp.png") + from src.utils import load_checkpoint + from models import UNetMNIST + + model = UNetMNIST(use_attention=args.attention).to(DEVICE) + if os.path.exists(args.checkpoint): + model = load_checkpoint(model, os.path.basename(args.checkpoint)) + + if args.forward or args.all: + visualize_forward_diffusion() + + if args.backward or args.all: + visualize_backward_diffusion(model) + + if args.all: + generate_grid(model) + + +if __name__ == "__main__": + main() diff --git a/models/__init__.py b/models/__init__.py new file mode 100644 index 0000000..db05aa9 --- /dev/null +++ b/models/__init__.py @@ -0,0 +1,6 @@ +""" +Model definitions for MNIST Diffusion. +""" + +from .unet import UNetMNIST +from .autoencoder import AEModule, Autoencoder diff --git a/models/autoencoder.py b/models/autoencoder.py new file mode 100644 index 0000000..353ec89 --- /dev/null +++ b/models/autoencoder.py @@ -0,0 +1,103 @@ +""" +Autoencoder model for MNIST. +Can be used for latent diffusion. +""" + +import torch +import torch.nn as nn + + +class AEModule(nn.Module): + """Base class for autoencoder modules (encoder and decoder).""" + + def __init__(self): + super().__init__() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + raise NotImplementedError("Subclasses must implement forward method.") + + def encode(self, x: torch.Tensor) -> torch.Tensor: + """Encode input to latent space.""" + raise NotImplementedError("Subclasses must implement encode method.") + + def decode(self, z: torch.Tensor) -> torch.Tensor: + """Decode from latent space to image space.""" + raise NotImplementedError("Subclasses must implement decode method.") + + +class Autoencoder(AEModule): + """ + Convolutional Autoencoder for MNIST images. + + Args: + input_channels: Number of input image channels (1 for MNIST) + latent_channels: Number of channels in latent space + """ + + def __init__(self, input_channels: int = 1, latent_channels: int = 1): + super().__init__() + self.input_channels = input_channels + self.latent_channels = latent_channels + + # Encoder: 28x28 -> 8x8 + self.encoder = nn.Sequential( + nn.Conv2d(input_channels, 16, kernel_size=3, padding=1), + nn.ReLU(), + nn.Conv2d(16, 16, kernel_size=3, padding=1), + nn.ReLU(), + nn.MaxPool2d(kernel_size=2), # 14x14 + + nn.Conv2d(16, 32, kernel_size=3, padding=1), + nn.ReLU(), + nn.Conv2d(32, 32, kernel_size=3, padding=1), + nn.ReLU(), + nn.MaxPool2d(kernel_size=2), # 7x7 + + nn.Conv2d(32, 64, kernel_size=3, padding=1), + nn.ReLU(), + nn.Conv2d(64, 64, kernel_size=3, padding=1), + nn.ReLU(), + nn.Conv2d(64, latent_channels, kernel_size=3, padding=1), + nn.ReLU(), + + # 7x7 -> 8x8 (no activation - latent space should be unconstrained) + nn.Conv2d(latent_channels, latent_channels, kernel_size=2, padding=1), + ) + + # Decoder: 8x8 -> 28x28 + self.decoder = nn.Sequential( + # 8x8 -> 7x7 + nn.Conv2d(latent_channels, latent_channels, kernel_size=2, padding=0), + nn.ReLU(), + + nn.Conv2d(latent_channels, 64, kernel_size=3, padding=1), + nn.ReLU(), + nn.Conv2d(64, 64, kernel_size=3, padding=1), + nn.ReLU(), + nn.ConvTranspose2d(64, 32, kernel_size=2, stride=2), # 14x14 + nn.ReLU(), + + nn.Conv2d(32, 32, kernel_size=3, padding=1), + nn.ReLU(), + nn.ConvTranspose2d(32, 16, kernel_size=2, stride=2), # 28x28 + nn.ReLU(), + + nn.Conv2d(16, 16, kernel_size=3, padding=1), + nn.ReLU(), + nn.Conv2d(16, input_channels, kernel_size=3, padding=1), + nn.Sigmoid(), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Full autoencoder forward pass.""" + z = self.encoder(x) + x_recon = self.decoder(z) + return x_recon + + def encode(self, x: torch.Tensor) -> torch.Tensor: + """Encode input to latent space.""" + return self.encoder(x) + + def decode(self, z: torch.Tensor) -> torch.Tensor: + """Decode from latent space to image space.""" + return self.decoder(z) diff --git a/models/unet.py b/models/unet.py new file mode 100644 index 0000000..f21859a --- /dev/null +++ b/models/unet.py @@ -0,0 +1,140 @@ +""" +UNet model for MNIST diffusion. +""" + +import torch +import torch.nn as nn +import torch.nn.functional as F + +import sys +sys.path.append("..") +from src.config import IMG_SIZE, NB_CHANNEL, NB_LABEL, DIFFU_STEPS + + +class SelfAttention(nn.Module): + """Self-attention module for UNet.""" + + def __init__(self, nb_channels: int, nb_heads: int): + super().__init__() + self.attention = nn.MultiheadAttention(nb_channels, nb_heads) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + _, c, w, h = x.shape + x = x.view(-1, c, w * h) + x = x.permute(2, 0, 1) + x, _ = self.attention(x, x, x) + x = x.permute(1, 2, 0) + x = x.view(-1, c, w, h) + return x + + +class UNetMNIST(nn.Module): + """ + UNet architecture for MNIST diffusion model. + + Inputs: + xt: image at step t (NB_CHANNEL x IMG_SIZE x IMG_SIZE) + t: step number (1) + vec: one-hot vector of the label (NB_LABEL) + + Output: + Predicted noise or denoised image (NB_CHANNEL x IMG_SIZE x IMG_SIZE) + """ + + def __init__(self, use_attention: bool = False): + super().__init__() + self.use_attention = use_attention + + # Encoder for timestep t + self.encodet = nn.Linear(1, IMG_SIZE * IMG_SIZE) + + # Encoder for label vector + self.encodevec = nn.Linear(NB_LABEL, IMG_SIZE * IMG_SIZE) + + # UNet encoder (2 extra channels for t and vec) + self.conv1 = nn.Conv2d(NB_CHANNEL + 2, 64, 3, padding=1) + self.conv2 = nn.Conv2d(64, 64, 3, padding=1) + self.maxpool1 = nn.MaxPool2d(2, 2) + + self.conv3 = nn.Conv2d(64, 128, 3, padding=1) + self.conv4 = nn.Conv2d(128, 128, 3, padding=1) + self.maxpool2 = nn.MaxPool2d(2, 2) + + # Bottleneck + self.conv5 = nn.Conv2d(128, 256, 3, padding=1) + self.conv6 = nn.Conv2d(256, 256, 3, padding=1) + + # Optional attention layers + if use_attention: + self.att1 = SelfAttention(128, 8) + self.att2 = SelfAttention(256, 8) + self.att3 = SelfAttention(256, 8) + + # UNet decoder + self.upconv1 = nn.ConvTranspose2d(256, 128, 2, stride=2) + self.conv7 = nn.Conv2d(256, 128, 3, padding=1) + self.conv8 = nn.Conv2d(128, 128, 3, padding=1) + + self.upconv2 = nn.ConvTranspose2d(128, 64, 2, stride=2) + self.conv9 = nn.Conv2d(128, 64, 3, padding=1) + self.conv10 = nn.Conv2d(64, 64, 3, padding=1) + + # Output layer + self.conv11 = nn.Conv2d(64, NB_CHANNEL, 3, padding=1) + + def forward(self, xt: torch.Tensor, t: torch.Tensor, vec: torch.Tensor) -> torch.Tensor: + """ + Forward pass of UNet. + + Args: + xt: Noisy image at timestep t [B, C, H, W] + t: Timestep [B, 1] + vec: Label one-hot vector [B, NB_LABEL] + + Returns: + Predicted noise or denoised image [B, C, H, W] + """ + # Encode timestep and label + t_enc = F.relu(self.encodet(t / DIFFU_STEPS)) + t_enc = t_enc.view(-1, 1, IMG_SIZE, IMG_SIZE) + + vec_enc = F.relu(self.encodevec(vec)) + vec_enc = vec_enc.view(-1, 1, IMG_SIZE, IMG_SIZE) + + # Concatenate input with embeddings + x = torch.cat((xt, t_enc, vec_enc), dim=1) + + # Encoder path + x1 = F.relu(self.conv1(x)) + x1 = F.relu(self.conv2(x1)) + + x2 = self.maxpool1(x1) + x2 = F.relu(self.conv3(x2)) + x2 = F.relu(self.conv4(x2)) + + # Bottleneck + x3 = self.maxpool2(x2) + if self.use_attention: + x3 = self.att1(x3) + x3 = F.relu(self.conv5(x3)) + if self.use_attention: + x3 = self.att2(x3) + x3 = F.relu(self.conv6(x3)) + if self.use_attention: + x3 = self.att3(x3) + + # Decoder path with skip connections + x4 = self.upconv1(x3) + x4 = torch.cat((x4, x2), dim=1) + x4 = F.relu(self.conv7(x4)) + x4 = F.relu(self.conv8(x4)) + + x5 = self.upconv2(x4) + x5 = torch.cat((x5, x1), dim=1) + x5 = F.relu(self.conv9(x5)) + x5 = F.relu(self.conv10(x5)) + + # Output + out = self.conv11(x5) + + return out diff --git a/plots.py b/plots.py deleted file mode 100644 index 0513495..0000000 --- a/plots.py +++ /dev/null @@ -1,174 +0,0 @@ -import matplotlib.pyplot as plt -from matplotlib.gridspec import GridSpec -import numpy as np -from torchvision import datasets -import torch -import os -from rich.progress import track - -from main import UNet, q_xt_xt_1, p_xt_1_xt, tensor_to_image -from autoencoder import Autoencoder - - -os.makedirs('plots', exist_ok=True) -os.makedirs('plots/diffusion', exist_ok=True) -os.makedirs('plots/diffusion_inverse', exist_ok=True) - - -DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu') - -DIFFU_STEPS = 1000 -BETA = torch.linspace(1e-4, 2e-2, DIFFU_STEPS, device=DEVICE) -BETA = torch.cat((torch.tensor([0.], device=DEVICE), BETA)) -ALPHA = 1 - BETA -ALPHA_BAR = torch.cumprod(ALPHA, dim=0) - -NB_BINS = 50 -BIN_MIN = -4 -BIN_MAX = 4 - - - -model = UNet().to(DEVICE) -model.load_state_dict(torch.load('model.pth')) - -autoencoder = Autoencoder(input_dim=(1, 28, 28), latent_dim=(1, 8, 8)).to(DEVICE) -autoencoder.load_state_dict(torch.load('autoencoder.pth')) - - - - -# plt.figure() -# plt.plot(BETA, label='beta') -# plt.plot(ALPHA, label='alpha') -# plt.plot(ALPHA_BAR, label='alpha_bar') -# plt.legend() -# plt.title('Alpha, Beta and Alpha_bar schedules') -# plt.savefig('plots/alpha_beta.tmp.png') - - -mnist = datasets.MNIST('data', train=True, download=True) -img, label = mnist[np.random.randint(0, len(mnist))] -img = np.array(img) / 255 * 2 - 1 -img = torch.tensor(img, device=DEVICE, dtype=torch.float32).unsqueeze(0).unsqueeze(0) - -encoded = autoencoder.encode(img) - -img = encoded.squeeze().cpu().detach().numpy() - -plt.figure() -plt.imshow(img, cmap='gray') -plt.title('Image') -plt.axis('off') -plt.savefig('plots/img.tmp.png') - - -plt.figure() -plt.hist(img.flatten(), bins=NB_BINS, range=(BIN_MIN, BIN_MAX)) -plt.yscale('log') -plt.title('Image histogram') -plt.savefig('plots/img_hist.tmp.png') - -def norm_dist(x, mean, std): - return np.exp(-0.5 * ((x - mean) / std) ** 2) / (std * np.sqrt(2 * np.pi)) - -x_norm = np.linspace(BIN_MIN, BIN_MAX, 100) -y_norm = norm_dist(x_norm, 0, 1) * 8**2 / NB_BINS * (BIN_MAX - BIN_MIN) - -fig = plt.figure(figsize=(10, 5)) -fig.suptitle('Diffusion naturelle') -gs = GridSpec(1, 3, figure=fig) -ax1 = fig.add_subplot(gs[0, 0]) -ax2 = fig.add_subplot(gs[0, 1:]) - -plots_to_save = np.linspace(1, DIFFU_STEPS, 100).astype(int) - -xt = torch.tensor(img, device=DEVICE, dtype=torch.float32).unsqueeze(0).unsqueeze(0) -for t in track(range(1, DIFFU_STEPS+1)): - xt, eps = q_xt_xt_1(xt, t) - - if t not in plots_to_save: - continue - - xt_numpy = xt.cpu().detach().numpy()[0, 0] - - ax1.clear() - ax1.imshow(xt_numpy, cmap='gray') - ax1.set_title(f'xt at t={t:04d}') - ax1.axis('off') - - ax2.clear() - ax2.hist(xt_numpy.flatten(), bins=NB_BINS, range=(BIN_MIN, BIN_MAX)) - ax2.plot(x_norm, y_norm, color='red', label='N(0, 1)') - ax2.legend() - ax2.set_yscale('log') - ax2.set_ylim(y_norm.min(), 1e3) - ax2.set_title(f'xt histogram') - - fig.savefig(f'plots/diffusion/{t:04d}.tmp.png') -os.system('convert -delay 20 -loop 0 plots/diffusion/*.png plots/diffusion.tmp.gif') - -fig = plt.figure(figsize=(10, 5)) -fig.suptitle('Diffusion inverse') -gs = GridSpec(1, 3, figure=fig) -ax1 = fig.add_subplot(gs[0, 0]) -ax2 = fig.add_subplot(gs[0, 1:]) - -xt = torch.randn(1, 1, 28, 28, device=DEVICE) -vec = torch.zeros(1, 10).to(DEVICE) -vec[0, label] = 1 -for t in track(range(DIFFU_STEPS, 0, -1)): - t_tensor = torch.tensor([[t]], device=DEVICE, dtype=torch.float32) - xt = p_xt_1_xt(model, xt, t_tensor, vec) - - if t not in plots_to_save: - continue - - xt_numpy = xt.cpu().detach().numpy()[0, 0] - - ax1.clear() - ax1.imshow(xt_numpy, cmap='gray') - ax1.set_title(f'xt at t={t:04d}') - ax1.axis('off') - - ax2.clear() - ax2.hist(xt_numpy.flatten(), bins=NB_BINS, range=(BIN_MIN, BIN_MAX)) - ax2.plot(x_norm, y_norm, color='red', label='N(0, 1)') - ax2.legend() - ax2.set_yscale('log') - ax2.set_ylim(y_norm.min(), 1e3) - ax2.set_title(f'xt histogram') - - fig.savefig(f'plots/diffusion_inverse/{t:04d}.tmp.png') -os.system('convert -delay 20 -loop 0 -reverse plots/diffusion_inverse/*.png plots/diffusion_inverse.tmp.gif') - -exit(0) - -tpause = {150: 'xt', 20: 'mu', 50: 'xt_1'} - -xT = xt = torch.randn(1, 1, 28, 28, device=DEVICE) -vec = torch.zeros(1, 10).to(DEVICE) -vec[0, label] = 1 - -for t in range(DIFFU_STEPS, 0, -1): - t_tensor = torch.tensor([[t]], device=DEVICE, dtype=torch.float32) - xt = p_xt_1_xt(model, xt, t_tensor, vec) - - if t in tpause: - plt.figure(figsize=(5, 5)) - plt.imshow(tensor_to_image(xt[0]), cmap='gray') - plt.axis('off') - plt.tight_layout() - plt.savefig(f'plots/{tpause[t]}.tmp.png') - -plt.figure(figsize=(5, 5)) -plt.imshow(tensor_to_image(xT[0]), cmap='gray') -plt.axis('off') -plt.tight_layout() -plt.savefig('plots/xT.tmp.png') - -plt.figure(figsize=(5, 5)) -plt.imshow(tensor_to_image(xt[0]), cmap='gray') -plt.axis('off') -plt.tight_layout() -plt.savefig('plots/x0.tmp.png') diff --git a/requirements.txt b/requirements.txt index d590e96..34490ec 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,5 +1,6 @@ -torchvision==0.15.2 -matplotlib==3.8.0 -rich==13.3.5 -PyQt5==5.15.10 -torch-trainer @ git+https://github.com/gdamms/torch-trainer.git@da70ca78adf6195f2e3add92938b9c282d5c94cb \ No newline at end of file +torch +torchvision +mlflow +scipy +matplotlib +rich diff --git a/src/__init__.py b/src/__init__.py new file mode 100644 index 0000000..220bb0e --- /dev/null +++ b/src/__init__.py @@ -0,0 +1,50 @@ +""" +MNIST Diffusion Model package. +""" + +from .config import ( + DEVICE, + IMG_SIZE, + NB_CHANNEL, + NB_LABEL, + DIFFU_STEPS, + ALPHA, + ALPHA_BAR, + BETA, + EPOCHS, + BATCH_SIZE, + LEARNING_RATE, + CHECKPOINT_DIR, + PLOTS_DIR, +) + +from .diffusion import ( + q_xt_xt_1, + q_xt_x0, + p_xt_1_xt, + p_xt_1_xt_x0_pred, + forward_diffusion, +) + +from .dataloader import ( + get_mnist_dataset, + get_diffusion_dataloader, + get_autoencoder_dataloader, + DiffusionDataset, + DiffusionDatasetX0, + AutoencoderDataset, +) + +from .utils import ( + tensor_to_image, + tensor_to_images, + save_checkpoint, + load_checkpoint, + save_plot, +) + +from .metrics import ( + fid, + kl_divergence, + jsd, +) diff --git a/src/config.py b/src/config.py new file mode 100644 index 0000000..4ac56c4 --- /dev/null +++ b/src/config.py @@ -0,0 +1,34 @@ +""" +Configuration file for MNIST Diffusion model. +Contains all hyperparameters and constants. +""" + +import torch + +# Device configuration +DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") + +# Image parameters (MNIST) +IMG_SIZE = 28 +NB_CHANNEL = 1 +NB_LABEL = 10 + +# Diffusion parameters +DIFFU_STEPS = 1000 + +# Noise schedule (linear beta schedule) +BETA = torch.linspace(1e-4, 2e-2, DIFFU_STEPS, device=DEVICE) +BETA = torch.cat((torch.tensor([0.0], device=DEVICE), BETA)) +ALPHA = 1 - BETA +ALPHA_BAR = torch.cumprod(ALPHA, dim=0) + +# Training parameters +EPOCHS = 10 +BATCH_SIZE = 64 +LEARNING_RATE = 2e-4 +NUM_WORKERS = 4 + +# Paths +DATA_DIR = "data" +CHECKPOINT_DIR = "checkpoints" +PLOTS_DIR = "plots" diff --git a/src/dataloader.py b/src/dataloader.py new file mode 100644 index 0000000..a08ae27 --- /dev/null +++ b/src/dataloader.py @@ -0,0 +1,210 @@ +""" +Data loading utilities for MNIST diffusion training. +""" + +import torch +from torch.utils.data import Dataset, DataLoader +from torchvision import datasets, transforms + +from .config import DEVICE, DIFFU_STEPS, NB_LABEL, DATA_DIR +from .diffusion import q_xt_x0 + + +def get_mnist_dataset(train: bool = True) -> datasets.MNIST: + """ + Load MNIST dataset. + + Args: + train: If True, load training set. Otherwise load test set. + + Returns: + MNIST dataset + """ + return datasets.MNIST( + root=DATA_DIR, + train=train, + download=True, + transform=transforms.ToTensor(), + ) + + +class DiffusionDataset(Dataset): + """ + Dataset wrapper for diffusion training. + Returns noisy image, timestep, label, and target noise. + + Args: + dataset: Base image dataset (e.g., MNIST) + autoencoder: Optional autoencoder for latent diffusion + """ + + def __init__(self, dataset: Dataset, autoencoder: torch.nn.Module = None): + super().__init__() + self.dataset = dataset + self.autoencoder = autoencoder + + def __getitem__(self, index: int) -> tuple[torch.Tensor, ...]: + # Get image and label + img, label = self.dataset[index] + img = img.to(DEVICE) + + # Optionally encode to latent space + if self.autoencoder is not None: + with torch.no_grad(): + img = self.autoencoder.encode(img.unsqueeze(0)).squeeze(0) + + # Normalize to [-1, 1] + img = img * 2 - 1 + + # Sample random timestep and add noise + t = torch.randint(1, DIFFU_STEPS + 1, (1,), device=DEVICE) + xt, eps = q_xt_x0(img, t) + + # Convert label to one-hot vector + vec = torch.nn.functional.one_hot( + torch.tensor(min(label, NB_LABEL - 1)), + num_classes=NB_LABEL, + ) + + return ( + xt.clone().detach().to(dtype=torch.float32, device=DEVICE), + t.clone().detach().to(dtype=torch.float32, device=DEVICE), + vec.clone().detach().to(dtype=torch.float32, device=DEVICE), + eps, # Target: the noise that was added + ) + + def __len__(self) -> int: + return len(self.dataset) + + +class DiffusionDatasetX0(Dataset): + """ + Dataset wrapper for diffusion training where model predicts x0 instead of noise. + Returns noisy image, timestep, label, and target clean image. + + Args: + dataset: Base image dataset (e.g., MNIST) + autoencoder: Optional autoencoder for latent diffusion + """ + + def __init__(self, dataset: Dataset, autoencoder: torch.nn.Module = None): + super().__init__() + self.dataset = dataset + self.autoencoder = autoencoder + + def __getitem__(self, index: int) -> tuple[torch.Tensor, ...]: + # Get image and label + img, label = self.dataset[index] + img = img.to(DEVICE) + + # Optionally encode to latent space + if self.autoencoder is not None: + with torch.no_grad(): + img = self.autoencoder.encode(img.unsqueeze(0)).squeeze(0) + + # Normalize to [-1, 1] + img = img * 2 - 1 + + # Sample random timestep and add noise + t = torch.randint(1, DIFFU_STEPS + 1, (1,), device=DEVICE) + xt, _ = q_xt_x0(img, t) + + # Convert label to one-hot vector + vec = torch.nn.functional.one_hot( + torch.tensor(min(label, NB_LABEL - 1)), + num_classes=NB_LABEL, + ) + + return ( + xt.clone().detach().to(dtype=torch.float32, device=DEVICE), + t.clone().detach().to(dtype=torch.float32, device=DEVICE), + vec.clone().detach().to(dtype=torch.float32, device=DEVICE), + img.clone().detach().to(dtype=torch.float32, device=DEVICE), # Target: clean image + ) + + def __len__(self) -> int: + return len(self.dataset) + + +class AutoencoderDataset(Dataset): + """ + Dataset wrapper for autoencoder training. + Returns image as both input and target. + + Args: + dataset: Base image dataset + """ + + def __init__(self, dataset: Dataset): + self.dataset = dataset + + def __len__(self) -> int: + return len(self.dataset) + + def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor]: + data = self.dataset[idx][0].to(DEVICE) + return data, data + + +def get_diffusion_dataloader( + predict_x0: bool = True, + batch_size: int = 64, + shuffle: bool = True, + num_workers: int = 4, + autoencoder: torch.nn.Module = None, +) -> DataLoader: + """ + Create a DataLoader for diffusion training. + + Args: + predict_x0: If True, model predicts x0. Otherwise predicts noise. + batch_size: Batch size + shuffle: Whether to shuffle data + num_workers: Number of data loading workers + autoencoder: Optional autoencoder for latent diffusion + + Returns: + DataLoader for training + """ + mnist = get_mnist_dataset(train=True) + + if predict_x0: + dataset = DiffusionDatasetX0(mnist, autoencoder) + else: + dataset = DiffusionDataset(mnist, autoencoder) + + return DataLoader( + dataset, + batch_size=batch_size, + shuffle=shuffle, + num_workers=num_workers, + persistent_workers=True if num_workers > 0 else False, + ) + + +def get_autoencoder_dataloader( + batch_size: int = 64, + shuffle: bool = True, + num_workers: int = 4, +) -> DataLoader: + """ + Create a DataLoader for autoencoder training. + + Args: + batch_size: Batch size + shuffle: Whether to shuffle data + num_workers: Number of data loading workers + + Returns: + DataLoader for training + """ + mnist = get_mnist_dataset(train=True) + dataset = AutoencoderDataset(mnist) + + return DataLoader( + dataset, + batch_size=batch_size, + shuffle=shuffle, + num_workers=num_workers, + persistent_workers=True if num_workers > 0 else False, + ) diff --git a/src/diffusion.py b/src/diffusion.py new file mode 100644 index 0000000..0d63ee1 --- /dev/null +++ b/src/diffusion.py @@ -0,0 +1,154 @@ +""" +Diffusion process utilities. +Contains forward and reverse diffusion functions. +""" + +import torch +from .config import DEVICE, ALPHA, ALPHA_BAR, BETA, DIFFU_STEPS + + +def q_xt_xt_1(xt_1: torch.Tensor, t: int | torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """ + Forward diffusion step: q(x_t | x_{t-1}). + Adds noise to image at step t-1 to get image at step t. + + Args: + xt_1: Image at timestep t-1 [B, C, H, W] + t: Timestep (int or tensor) + + Returns: + xt: Noisy image at timestep t + eps: The noise that was added + """ + if isinstance(t, int): + t_ind = torch.tensor(t, dtype=torch.long, device=DEVICE) + else: + t_ind = t.to(dtype=torch.long, device=DEVICE) + + alpha = ALPHA[t_ind] + mean = torch.sqrt(alpha) * xt_1 + std = torch.sqrt(1 - alpha) + + eps = torch.randn(xt_1.shape, device=DEVICE) + xt = mean + std * eps + + return xt, eps + + +def q_xt_x0(x0: torch.Tensor, t: int | torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """ + Forward diffusion: q(x_t | x_0). + Directly compute noisy image at any timestep t from clean image x0. + + Args: + x0: Clean image [B, C, H, W] or [C, H, W] + t: Timestep tensor [B, 1] or [B] + + Returns: + xt: Noisy image at timestep t + eps: The noise that was added + """ + if isinstance(t, int): + t_ind = torch.tensor(t, dtype=torch.long, device=DEVICE) + else: + t_ind = t.to(dtype=torch.long, device=DEVICE) + + # Handle both batched and single images + reshaped = len(x0.shape) == 3 + if reshaped: + c, w, h = x0.shape + b = 1 + x0 = x0.view(b, c, w, h) + else: + b, c, w, h = x0.shape + + # Reshape t for broadcasting + t_ind = t_ind.view(b, 1, 1, 1) + t_ind = t_ind.expand(b, c, w, h) + + alpha_bar = ALPHA_BAR[t_ind] + mean = torch.sqrt(alpha_bar) * x0 + std = torch.sqrt(1 - alpha_bar) + + eps = torch.randn(x0.shape, device=DEVICE) + xt = mean + std * eps + + if reshaped: + xt = xt.view(c, w, h) + + return xt, eps + + +def p_xt_1_xt(model: torch.nn.Module, xt: torch.Tensor, t: torch.Tensor, + vec: torch.Tensor) -> torch.Tensor: + """ + Reverse diffusion step: p(x_{t-1} | x_t). + Denoise image at step t to get image at step t-1. + Model predicts the noise. + + Args: + model: UNet model that predicts noise + xt: Noisy image at timestep t [B, C, H, W] + t: Timestep tensor [B, 1] + vec: Label one-hot vector [B, NB_LABEL] + + Returns: + xt_1: Denoised image at timestep t-1 + """ + t_ind = t.to(dtype=torch.long) if isinstance(t, torch.Tensor) else t + + alpha_bar_t = ALPHA_BAR[t_ind].view(-1, 1, 1, 1) + alpha_bar_t_1 = ALPHA_BAR[t_ind - 1].view(-1, 1, 1, 1) + alpha_t = ALPHA[t_ind].view(-1, 1, 1, 1) + beta_t = BETA[t_ind].view(-1, 1, 1, 1) + + beta_tilde = (1 - alpha_bar_t_1) / (1 - alpha_bar_t) * beta_t + + # Model predicts the noise + epsilon_theta = model(xt, t, vec) + + sigma_theta = torch.sqrt(beta_tilde) + mu_theta = (xt - (1 - alpha_t) / torch.sqrt(1 - alpha_bar_t) * epsilon_theta) / torch.sqrt(alpha_t) + + # Don't add noise at t=1 + mask_t0 = (t > 1).to(dtype=torch.float32).view(-1, 1, 1, 1) + noise = torch.randn(xt.shape, device=DEVICE) * mask_t0 + + return mu_theta + sigma_theta * noise + + +def p_xt_1_xt_x0_pred(model: torch.nn.Module, xt: torch.Tensor, t: torch.Tensor, + vec: torch.Tensor) -> torch.Tensor: + """ + Reverse diffusion step where model predicts x0 directly. + + Args: + model: UNet model that predicts clean image x0 + xt: Noisy image at timestep t [B, C, H, W] + t: Timestep tensor [B, 1] + vec: Label one-hot vector [B, NB_LABEL] + + Returns: + xt_1: Denoised image at timestep t-1 + """ + x0_pred = model(xt, t, vec) + xt_1, _ = q_xt_x0(x0_pred, t - 1) + return xt_1 + + +def forward_diffusion(x0: torch.Tensor) -> list[torch.Tensor]: + """ + Run full forward diffusion process. + + Args: + x0: Clean image [C, H, W] or [B, C, H, W] + + Returns: + List of images at each timestep [x0, x1, ..., xT] + """ + x = x0.clone() + xs = [x] + for t in range(1, DIFFU_STEPS + 1): + x, _ = q_xt_xt_1(x, t) + xs.append(x) + return xs diff --git a/src/metrics.py b/src/metrics.py new file mode 100644 index 0000000..08eb3cd --- /dev/null +++ b/src/metrics.py @@ -0,0 +1,98 @@ +import numpy as np +import scipy + + +def fid(reals: np.ndarray, fakes: np.ndarray) -> float: + """ + Calculate Frechet Inception Distance (FID) score. + + Args: + reals: Real images [N, C, H, W] + fakes: Generated images [N, C, H, W] + + Returns: + FID score (lower is better) + """ + reals = reals.reshape(reals.shape[0], -1) + fakes = fakes.reshape(fakes.shape[0], -1) + + mu_real = np.mean(reals, axis=0) + mu_fake = np.mean(fakes, axis=0) + sigma_real = np.cov(reals, rowvar=False) + sigma_fake = np.cov(fakes, rowvar=False) + + diff = mu_real - mu_fake + covmean, _ = scipy.linalg.sqrtm(sigma_real.dot(sigma_fake), disp=False) + + if not np.isfinite(covmean).all(): + eps = 1e-6 + offset = np.eye(sigma_real.shape[0]) * eps + covmean = scipy.linalg.sqrtm((sigma_real + offset).dot(sigma_fake + offset)) + + if np.iscomplexobj(covmean): + covmean = covmean.real + + return diff @ diff + np.trace(sigma_real) + np.trace(sigma_fake) - 2 * np.trace(covmean) + + +def kl_divergence(reals: np.ndarray, fakes: np.ndarray) -> float: + """ + Calculate KL divergence between real and fake image distributions. + + Args: + reals: Real images [N, C, H, W] + fakes: Generated images [N, C, H, W] + + Returns: + KL divergence value + """ + reals = reals.transpose(1, 0, 2, 3).reshape(reals.shape[1], -1) + fakes = fakes.transpose(1, 0, 2, 3).reshape(fakes.shape[1], -1) + + hist_real = np.apply_along_axis( + lambda a: np.histogram(a, bins=40, range=(-1, 1))[0], 1, reals + ) + hist_fake = np.apply_along_axis( + lambda a: np.histogram(a, bins=40, range=(-1, 1))[0], 1, fakes + ) + + # Add smoothing + hist_real = hist_real + 1 + hist_fake = hist_fake + 1 + + hist_real = hist_real / np.sum(hist_real) + hist_fake = hist_fake / np.sum(hist_fake) + + return np.mean(np.log(hist_real / hist_fake)) + + +def jsd(reals: np.ndarray, fakes: np.ndarray) -> float: + """ + Calculate Jensen-Shannon divergence between real and fake image distributions. + + Args: + reals: Real images [N, C, H, W] + fakes: Generated images [N, C, H, W] + + Returns: + JSD value + """ + reals = reals.transpose(1, 0, 2, 3).reshape(reals.shape[1], -1) + fakes = fakes.transpose(1, 0, 2, 3).reshape(fakes.shape[1], -1) + + hist_real = np.apply_along_axis( + lambda a: np.histogram(a, bins=40, range=(-1, 1))[0], 1, reals + ) + hist_fake = np.apply_along_axis( + lambda a: np.histogram(a, bins=40, range=(-1, 1))[0], 1, fakes + ) + + hist_real = hist_real + 1 + hist_fake = hist_fake + 1 + + hist_real = hist_real / np.sum(hist_real) + hist_fake = hist_fake / np.sum(hist_fake) + + hist_avg = (hist_real + hist_fake) / 2 + + return 0.5 * (np.mean(np.log(hist_real / hist_avg)) + np.mean(np.log(hist_fake / hist_avg))) diff --git a/src/sample.py b/src/sample.py new file mode 100644 index 0000000..5b438ed --- /dev/null +++ b/src/sample.py @@ -0,0 +1,248 @@ +""" +Sampling and evaluation script for trained diffusion model. +Generate samples from a trained model and visualize results. +""" + +from models import UNetMNIST +from src.utils import ensure_dirs, tensor_to_image, load_checkpoint +from src.dataloader import get_mnist_dataset +from src.diffusion import p_xt_1_xt_x0_pred, forward_diffusion, q_xt_x0 +from src.config import ( + DEVICE, DIFFU_STEPS, NB_CHANNEL, IMG_SIZE, NB_LABEL, + CHECKPOINT_DIR, PLOTS_DIR +) +import os +import torch +import numpy as np +import matplotlib.pyplot as plt +from rich.progress import track + +import sys +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + + +def generate_samples( + model: torch.nn.Module, + n_samples: int = 10, + labels: list[int] | None = None, +) -> torch.Tensor: + """ + Generate samples from the diffusion model. + + Args: + model: Trained UNet model + n_samples: Number of samples per class (or total if labels provided) + labels: Optional list of specific labels to generate + + Returns: + Generated samples tensor [N, C, H, W] + """ + model.eval() + + if labels is None: + # Generate samples for all classes + labels = list(range(NB_LABEL)) * n_samples + + n_total = len(labels) + + with torch.no_grad(): + # Start from pure noise + x = torch.randn(n_total, NB_CHANNEL, IMG_SIZE, IMG_SIZE, device=DEVICE) + + # Create one-hot label vectors + vec = torch.tensor(labels, device=DEVICE) + vec = torch.nn.functional.one_hot(vec, num_classes=NB_LABEL).to(dtype=torch.float32) + + # Reverse diffusion process + for t in track(range(DIFFU_STEPS, 0, -1), description="Generating samples"): + t_tensor = torch.tensor([[t]] * n_total, device=DEVICE, dtype=torch.float32) + x = p_xt_1_xt_x0_pred(model, x, t_tensor, vec) + + # Normalize to [0, 1] + x = x * 0.5 + 0.5 + x = x.clamp(0, 1) + + return x + + +def visualize_forward_diffusion(save_path: str | None = None): + """ + Visualize the forward diffusion process on a real image. + + Args: + save_path: Path to save the visualization + """ + ensure_dirs() + + # Get a random image from MNIST + dataset = get_mnist_dataset(train=True) + idx = np.random.randint(0, len(dataset)) + img, label = dataset[idx] + img = img.to(DEVICE) + + # Normalize to [-1, 1] + img = img * 2 - 1 + + # Run forward diffusion + xs = forward_diffusion(img) + + # Select timesteps to visualize + n_plots = 10 + timesteps = np.linspace(1, DIFFU_STEPS, n_plots, dtype=int) + + fig, axes = plt.subplots(2, n_plots + 1, figsize=(2 * n_plots, 5)) + + # Row labels + axes[0, 0].text(0.5, 0.5, 'Step-by-step', ha='center', va='center', fontsize=10) + axes[0, 0].axis('off') + axes[1, 0].text(0.5, 0.5, 'Direct', ha='center', va='center', fontsize=10) + axes[1, 0].axis('off') + + # Plot step-by-step diffusion + for i, t in enumerate(timesteps): + axes[0, i + 1].imshow(tensor_to_image(xs[t]), cmap='gray') + axes[0, i + 1].set_title(f't={t}') + axes[0, i + 1].axis('off') + + # Plot direct diffusion for comparison + xt, _ = q_xt_x0(img, t) + axes[1, i + 1].imshow(tensor_to_image(xt), cmap='gray') + axes[1, i + 1].axis('off') + + fig.suptitle(f'Forward Diffusion Process (Label: {label})') + plt.tight_layout() + + if save_path is None: + save_path = os.path.join(PLOTS_DIR, 'forward_diffusion.png') + fig.savefig(save_path) + print(f"Saved forward diffusion visualization to {save_path}") + plt.close(fig) + + +def visualize_backward_diffusion(model: torch.nn.Module, save_path: str | None = None): + """ + Visualize the backward (reverse) diffusion process. + + Args: + model: Trained UNet model + save_path: Path to save the visualization + """ + ensure_dirs() + model.eval() + + n_classes = NB_LABEL + n_timesteps = 10 + timesteps = np.linspace(1, DIFFU_STEPS, n_timesteps, dtype=int)[::-1] + + fig, axes = plt.subplots(n_classes, n_timesteps, figsize=(2 * n_timesteps, 2 * n_classes)) + + with torch.no_grad(): + # Start from noise + x = torch.randn(n_classes, NB_CHANNEL, IMG_SIZE, IMG_SIZE, device=DEVICE) + + # One sample per class + vec = torch.arange(n_classes, device=DEVICE) + vec = torch.nn.functional.one_hot(vec, num_classes=NB_LABEL).to(dtype=torch.float32) + + for t in track(range(DIFFU_STEPS, 0, -1), description="Visualizing backward diffusion"): + t_tensor = torch.tensor([[t]] * n_classes, device=DEVICE, dtype=torch.float32) + x = p_xt_1_xt_x0_pred(model, x, t_tensor, vec) + + if t in timesteps: + t_idx = timesteps.tolist().index(t) + for class_idx in range(n_classes): + axes[class_idx, t_idx].imshow(tensor_to_image(x[class_idx]), cmap='gray') + if class_idx == 0: + axes[class_idx, t_idx].set_title(f't={t}') + if t_idx == 0: + axes[class_idx, t_idx].set_ylabel(f'Class {class_idx}') + axes[class_idx, t_idx].set_xticks([]) + axes[class_idx, t_idx].set_yticks([]) + + fig.suptitle('Backward Diffusion Process') + plt.tight_layout() + + if save_path is None: + save_path = os.path.join(PLOTS_DIR, 'backward_diffusion.png') + fig.savefig(save_path) + print(f"Saved backward diffusion visualization to {save_path}") + plt.close(fig) + + +def generate_grid(model: torch.nn.Module, n_per_class: int = 10, save_path: str | None = None): + """ + Generate a grid of samples, organized by class. + + Args: + model: Trained UNet model + n_per_class: Number of samples per class + save_path: Path to save the grid + """ + ensure_dirs() + + # Generate samples + labels = [] + for class_idx in range(NB_LABEL): + labels.extend([class_idx] * n_per_class) + + samples = generate_samples(model, labels=labels) + samples = samples.cpu().numpy() + + # Create grid + fig, axes = plt.subplots(NB_LABEL, n_per_class, figsize=(n_per_class, NB_LABEL)) + + for class_idx in range(NB_LABEL): + for sample_idx in range(n_per_class): + idx = class_idx * n_per_class + sample_idx + axes[class_idx, sample_idx].imshow(samples[idx].transpose(1, 2, 0).squeeze(), cmap='gray') + axes[class_idx, sample_idx].axis('off') + + if sample_idx == 0: + axes[class_idx, sample_idx].set_ylabel(f'{class_idx}') + + fig.suptitle('Generated MNIST Digits') + plt.tight_layout() + + if save_path is None: + save_path = os.path.join(PLOTS_DIR, 'generated_grid.png') + fig.savefig(save_path) + print(f"Saved generated grid to {save_path}") + plt.close(fig) + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser(description="Sample from trained diffusion model") + parser.add_argument("--checkpoint", type=str, default=os.path.join(CHECKPOINT_DIR, "diffusion_latest.pt"), + help="Path to model checkpoint") + parser.add_argument("--n-samples", type=int, default=10, help="Samples per class") + parser.add_argument("--attention", action="store_true", help="Use attention in model") + parser.add_argument("--forward", action="store_true", help="Visualize forward diffusion") + parser.add_argument("--backward", action="store_true", help="Visualize backward diffusion") + parser.add_argument("--grid", action="store_true", help="Generate sample grid") + parser.add_argument("--all", action="store_true", help="Run all visualizations") + + args = parser.parse_args() + + # Load model + model = UNetMNIST(use_attention=args.attention).to(DEVICE) + + if os.path.exists(args.checkpoint): + model = load_checkpoint(model, os.path.basename(args.checkpoint)) + else: + print(f"Warning: Checkpoint {args.checkpoint} not found. Using untrained model.") + + # Run visualizations + if args.forward or args.all: + visualize_forward_diffusion() + + if args.backward or args.all: + visualize_backward_diffusion(model) + + if args.grid or args.all: + generate_grid(model, n_per_class=args.n_samples) + + # Default: generate grid if no specific option selected + if not (args.forward or args.backward or args.grid or args.all): + generate_grid(model, n_per_class=args.n_samples) diff --git a/src/train_autoencoder.py b/src/train_autoencoder.py new file mode 100644 index 0000000..669ce47 --- /dev/null +++ b/src/train_autoencoder.py @@ -0,0 +1,167 @@ +""" +Training script for MNIST autoencoder. +""" + +from models import Autoencoder, AEModule +from src.utils import ensure_dirs, save_checkpoint +from src.dataloader import get_autoencoder_dataloader +from src.config import DEVICE, BATCH_SIZE, NUM_WORKERS, CHECKPOINT_DIR, PLOTS_DIR +import os +import torch +import torch.nn as nn +import matplotlib.pyplot as plt +from rich.progress import track +import mlflow + +import sys +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + + +def train_autoencoder( + epochs: int = 10, + learning_rate: float = 1e-3, + batch_size: int = BATCH_SIZE, + latent_channels: int = 1, + checkpoint_path: str | None = None, + run_name: str | None = None, +): + """ + Train the autoencoder model. + + Args: + epochs: Number of training epochs + learning_rate: Learning rate for optimizer + batch_size: Training batch size + latent_channels: Number of channels in latent space + checkpoint_path: Path to checkpoint to resume training from + run_name: Name for this training run (for logging) + """ + ensure_dirs() + + # Initialize model + model = Autoencoder(input_channels=1, latent_channels=latent_channels).to(DEVICE) + + if checkpoint_path: + if not os.path.exists(checkpoint_path): + print(f"Checkpoint not found: {checkpoint_path}") + return + model.load_state_dict(torch.load(checkpoint_path, weights_only=True)) + + if run_name is None: + from datetime import datetime + run_name = f"autoencoder_{datetime.now().strftime('%Y%m%d_%H%M%S')}" + + mlflow.set_experiment("MNIST Autoencoder") + mlflow.start_run(run_name=run_name) + + # Setup training + optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate) + # Use MSE loss instead of BCE for better reconstruction of continuous values + criterion = nn.MSELoss() + + # Get dataloader + dataloader = get_autoencoder_dataloader( + batch_size=batch_size, + num_workers=NUM_WORKERS, + ) + + # Training loop + for epoch in range(1, epochs + 1): + model.train() + epoch_loss = 0.0 + + for batch_idx, (x, target) in enumerate(track(dataloader, description=f"Epoch {epoch}/{epochs}")): + optimizer.zero_grad() + + # Forward pass + x_recon = model(x) + loss = criterion(x_recon, target) + + # Backward pass + loss.backward() + + # Gradient clipping to prevent exploding gradients + torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) + + optimizer.step() + + epoch_loss += loss.item() + + avg_loss = epoch_loss / len(dataloader) + mlflow.log_metric("epoch_loss", avg_loss, step=epoch) + + # Save checkpoint + save_checkpoint(model, f"autoencoder_epoch_{epoch:03d}.pt") + save_checkpoint(model, "autoencoder_latest.pt") + + # Visualize results + visualize_reconstructions(model, dataloader) + + mlflow.end_run() + return model + + +def visualize_reconstructions(model: AEModule, dataloader, n_samples: int = 10): + """Visualize original, latent, and reconstructed images.""" + model.eval() + ensure_dirs() + + # Get a batch of samples + x_batch, _ = next(iter(dataloader)) + x_batch = x_batch[:n_samples] + + with torch.no_grad(): + latent = model.encode(x_batch) + x_recon = model.decode(latent) + + # Create visualization + fig, axes = plt.subplots(3, n_samples + 1, figsize=(2 * n_samples, 6)) + + # Labels + axes[0, 0].text(0.5, 0.5, 'Original', ha='center', va='center', fontsize=12) + axes[0, 0].axis('off') + axes[1, 0].text(0.5, 0.5, 'Latent', ha='center', va='center', fontsize=12) + axes[1, 0].axis('off') + axes[2, 0].text(0.5, 0.5, 'Reconstructed', ha='center', va='center', fontsize=12) + axes[2, 0].axis('off') + + # Plot images + for i in range(n_samples): + axes[0, i + 1].imshow(x_batch[i].cpu().squeeze().numpy(), cmap='gray') + axes[0, i + 1].axis('off') + + axes[1, i + 1].imshow(latent[i].cpu().squeeze().numpy(), cmap='gray') + axes[1, i + 1].axis('off') + + axes[2, i + 1].imshow(x_recon[i].cpu().squeeze().numpy(), cmap='gray') + axes[2, i + 1].axis('off') + + fig.suptitle('Autoencoder Results') + plt.tight_layout() + + save_path = os.path.join(PLOTS_DIR, 'autoencoder_results.png') + fig.savefig(save_path) + plt.close(fig) + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser(description="Train MNIST autoencoder") + parser.add_argument("--epochs", type=int, default=10, help="Number of epochs") + parser.add_argument("--lr", type=float, default=1e-3, help="Learning rate") + parser.add_argument("--batch-size", type=int, default=BATCH_SIZE, help="Batch size") + parser.add_argument("--latent-channels", type=int, default=1, help="Latent channels") + parser.add_argument("--checkpoint", type=str, default=None, help="Resume from checkpoint") + + args = parser.parse_args() + + torch.multiprocessing.set_start_method("spawn", force=True) + + train_autoencoder( + epochs=args.epochs, + learning_rate=args.lr, + batch_size=args.batch_size, + latent_channels=args.latent_channels, + checkpoint_path=args.checkpoint, + ) diff --git a/src/train_diffusion.py b/src/train_diffusion.py new file mode 100644 index 0000000..cb70650 --- /dev/null +++ b/src/train_diffusion.py @@ -0,0 +1,197 @@ +""" +Training script for MNIST diffusion model. +""" + +from models import UNetMNIST +from src.utils import ( + ensure_dirs, save_checkpoint, tensor_to_image, + figure_to_image, +) +from src.metrics import fid, kl_divergence, jsd +from src.diffusion import p_xt_1_xt_x0_pred +from src.dataloader import get_diffusion_dataloader, get_mnist_dataset +from src.config import ( + DEVICE, EPOCHS, LEARNING_RATE, BATCH_SIZE, NUM_WORKERS, + DIFFU_STEPS, NB_CHANNEL, IMG_SIZE, NB_LABEL, CHECKPOINT_DIR, PLOTS_DIR +) +import os +import torch +import torch.nn as nn +import numpy as np +import matplotlib.pyplot as plt +from rich.progress import track +import mlflow + +import sys +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + + +def train_diffusion( + epochs: int = EPOCHS, + learning_rate: float = LEARNING_RATE, + batch_size: int = BATCH_SIZE, + predict_x0: bool = True, + use_attention: bool = False, + checkpoint_path: str | None = None, + run_name: str | None = None, +): + """ + Train the diffusion model. + + Args: + epochs: Number of training epochs + learning_rate: Learning rate for optimizer + batch_size: Training batch size + predict_x0: If True, model predicts x0. Otherwise predicts noise. + use_attention: If True, use self-attention in UNet + checkpoint_path: Path to checkpoint to resume training from + run_name: Name for this training run (for logging) + """ + ensure_dirs() + + # Setup run name and logging + if run_name is None: + from datetime import datetime + run_name = f"diffusion_{datetime.now().strftime('%Y%m%d_%H%M%S')}" + + mlflow.set_experiment("MNIST Diffusion") + mlflow.start_run(run_name=run_name) + + # Initialize model + model = UNetMNIST(use_attention=use_attention).to(DEVICE) + + if checkpoint_path and os.path.exists(checkpoint_path): + model.load_state_dict(torch.load(checkpoint_path, weights_only=True)) + print(f"Loaded checkpoint from {checkpoint_path}") + + # Setup training + optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate) + criterion = nn.MSELoss() + + # Get dataloader + dataloader = get_diffusion_dataloader( + predict_x0=predict_x0, + batch_size=batch_size, + num_workers=NUM_WORKERS, + ) + + # Training loop + global_step = 0 + for epoch in range(1, epochs + 1): + model.train() + epoch_loss = 0.0 + + for batch_idx, (xt, t, vec, target) in enumerate(track(dataloader, description=f"Epoch {epoch}/{epochs}")): + optimizer.zero_grad() + + # Forward pass + pred = model(xt, t, vec) + loss = criterion(pred, target) + + # Backward pass + loss.backward() + optimizer.step() + + epoch_loss += loss.item() + global_step += 1 + + # Log training loss every 100 steps + if global_step % 100 == 0: + avg_loss = epoch_loss / (batch_idx + 1) + mlflow.log_metric("train_loss", avg_loss, step=global_step) + + avg_loss = epoch_loss / len(dataloader) + mlflow.log_metric("epoch_loss", avg_loss, step=epoch) + + # Save checkpoint every epoch + save_checkpoint(model, f"diffusion_epoch_{epoch:03d}.pt") + save_checkpoint(model, "diffusion_latest.pt") + + # Generate and log sample images + if epoch % 1 == 0: + evaluate_and_log(model, epoch, predict_x0) + + mlflow.end_run() + return model + + +def evaluate_and_log(model: nn.Module, epoch: int, predict_x0: bool = True): + """Generate samples and log metrics.""" + model.eval() + + with torch.no_grad(): + # Generate samples + batch_size = 64 + n_batches = 4 + + fakes = [] + for _ in range(n_batches): + x = torch.randn(batch_size, NB_CHANNEL, IMG_SIZE, IMG_SIZE, device=DEVICE) + vec = torch.randint(0, NB_LABEL, (batch_size,), device=DEVICE) + vec = torch.nn.functional.one_hot(vec, num_classes=NB_LABEL).to(dtype=torch.float32) + + # Reverse diffusion + for t in range(DIFFU_STEPS, 0, -1): + t_tensor = torch.tensor([[t]] * batch_size, device=DEVICE, dtype=torch.float32) + x = p_xt_1_xt_x0_pred(model, x, t_tensor, vec) + + x = x.cpu().numpy() + x = (x - x.min()) / (x.max() - x.min() + 1e-8) + fakes.append(x) + + fakes = np.concatenate(fakes) + + # Get real samples for comparison + dataset = get_mnist_dataset(train=True) + n_samples = len(fakes) + reals = torch.stack([dataset[i][0] for i in range(n_samples)]).numpy() + reals = reals * 2 - 1 + + # Log metrics + fid_score = fid(reals, fakes) + kl_score = kl_divergence(reals, fakes) + jsd_score = jsd(reals, fakes) + + mlflow.log_metric("FID", fid_score, step=epoch) + mlflow.log_metric("KL Divergence", kl_score, step=epoch) + mlflow.log_metric("JSD", jsd_score, step=epoch) + + # Log sample images + fig, axes = plt.subplots(4, 8, figsize=(16, 8)) + for i, ax in enumerate(axes.flat): + if i < len(fakes): + ax.imshow(fakes[i].transpose(1, 2, 0).squeeze(), cmap='gray') + ax.axis('off') + fig.suptitle(f"Generated Samples - Epoch {epoch}") + plt.tight_layout() + + mlflow.log_figure(fig, f"samples_epoch_{epoch:03d}.png") + + # Save to plots folder + fig.savefig(os.path.join(PLOTS_DIR, f"samples_epoch_{epoch:03d}.png")) + plt.close(fig) + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser(description="Train MNIST diffusion model") + parser.add_argument("--epochs", type=int, default=EPOCHS, help="Number of epochs") + parser.add_argument("--lr", type=float, default=LEARNING_RATE, help="Learning rate") + parser.add_argument("--batch-size", type=int, default=BATCH_SIZE, help="Batch size") + parser.add_argument("--attention", action="store_true", help="Use self-attention") + parser.add_argument("--checkpoint", type=str, default=None, help="Resume from checkpoint") + parser.add_argument("--name", type=str, default=None, help="Run name") + + args = parser.parse_args() + + torch.multiprocessing.set_start_method("spawn", force=True) + + train_diffusion( + epochs=args.epochs, + learning_rate=args.lr, + batch_size=args.batch_size, + use_attention=args.attention, + checkpoint_path=args.checkpoint, + run_name=args.name, + ) diff --git a/src/utils.py b/src/utils.py new file mode 100644 index 0000000..f727a68 --- /dev/null +++ b/src/utils.py @@ -0,0 +1,110 @@ +""" +Utility functions for MNIST diffusion. +""" + +import os +import io +import numpy as np +import torch +import scipy.linalg +from PIL import Image +import matplotlib.pyplot as plt + +from .config import CHECKPOINT_DIR, PLOTS_DIR + + +def ensure_dirs(): + """Create necessary directories if they don't exist.""" + os.makedirs(CHECKPOINT_DIR, exist_ok=True) + os.makedirs(PLOTS_DIR, exist_ok=True) + + +def tensor_to_image(tensor: torch.Tensor) -> np.ndarray: + """ + Convert a tensor to a numpy image array. + + Args: + tensor: Image tensor [C, H, W] + + Returns: + Numpy array [H, W, C] normalized to [0, 1] + """ + img = tensor.clone().detach().cpu().numpy().transpose(1, 2, 0) + img -= img.min() + img /= img.max() + 1e-8 + return img + + +def tensor_to_images(tensor: torch.Tensor) -> np.ndarray: + """ + Convert a batch of tensors to numpy image arrays. + + Args: + tensor: Batch of image tensors [B, C, H, W] + + Returns: + Numpy array [B, H, W, C] normalized to [0, 1] + """ + img = tensor.clone().detach().cpu().numpy().transpose(0, 2, 3, 1) + img -= np.min(img, axis=(1, 2, 3), keepdims=True) + img /= np.max(img, axis=(1, 2, 3), keepdims=True) + 1e-8 + return img + + +def figure_to_image(figure: plt.Figure) -> np.ndarray: + """ + Convert a matplotlib figure to a numpy image array. + + Args: + figure: Matplotlib figure + + Returns: + Numpy array of the figure image + """ + buf = io.BytesIO() + figure.savefig(buf, format='png') + buf.seek(0) + image = np.array(Image.open(buf)) + return image + + +def save_checkpoint(model: torch.nn.Module, filename: str): + """ + Save model checkpoint. + + Args: + model: Model to save + filename: Filename (will be saved in CHECKPOINT_DIR) + """ + ensure_dirs() + path = os.path.join(CHECKPOINT_DIR, filename) + torch.save(model.state_dict(), path) + + +def load_checkpoint(model: torch.nn.Module, filename: str) -> torch.nn.Module: + """ + Load model checkpoint. + + Args: + model: Model architecture to load weights into + filename: Filename (loaded from CHECKPOINT_DIR) + + Returns: + Model with loaded weights + """ + path = os.path.join(CHECKPOINT_DIR, filename) + model.load_state_dict(torch.load(path, weights_only=True)) + return model + + +def save_plot(figure: plt.Figure, filename: str): + """ + Save a matplotlib figure to the plots directory. + + Args: + figure: Matplotlib figure to save + filename: Filename (will be saved in PLOTS_DIR) + """ + ensure_dirs() + path = os.path.join(PLOTS_DIR, filename) + figure.savefig(path) diff --git a/utils.py b/utils.py deleted file mode 100644 index 9935d16..0000000 --- a/utils.py +++ /dev/null @@ -1,108 +0,0 @@ -import numpy as np -import scipy.linalg -import cv2 - - -def fid(reals, fakes): - """FID score calculation. - - Args: - reals (numpy.array): Real images. - fakes (numpy.array): Fake images. - """ - reals = reals.reshape(reals.shape[0], -1) - fakes = fakes.reshape(fakes.shape[0], -1) - - mu_real = np.mean(reals, axis=0) - mu_fake = np.mean(fakes, axis=0) - sigma_real = np.cov(reals, rowvar=False) - sigma_fake = np.cov(fakes, rowvar=False) - - diff = mu_real - mu_fake - covmean = np.dot(sigma_real, sigma_fake.T) - covmean, _ = scipy.linalg.sqrtm(sigma_real.dot(sigma_fake), disp=False) - - if not np.isfinite(covmean).all(): - eps=1e-6 - offset = np.eye(sigma_real.shape[0]) * eps - ncovmean = scipy.linalg.sqrtm((sigma_real + offset).dot(sigma_fake + offset)) - covmean = ncovmean - - if np.iscomplexobj(covmean): - covmean = covmean.real - - return diff @ diff + np.trace(sigma_real) + np.trace(sigma_fake) - 2 * np.trace(covmean) - - -def kl(reals, fakes): - """KL divergence calculation. - - Args: - reals (numpy.array): Real images. - fakes (numpy.array): Fake images. - """ - reals = reals.transpose(1, 0, 2, 3).reshape(reals.shape[1], -1) - fakes = fakes.transpose(1, 0, 2, 3).reshape(fakes.shape[1], -1) - - hist_real = np.apply_along_axis(lambda a: np.histogram(a, bins=40, range=(-1, 1))[0], 1, reals) - hist_fake = np.apply_along_axis(lambda a: np.histogram(a, bins=40, range=(-1, 1))[0], 1, fakes) - - hist_real = hist_real + 1 - hist_fake = hist_fake + 1 - - hist_real = hist_real / np.sum(hist_real) - hist_fake = hist_fake / np.sum(hist_fake) - - return np.mean(np.log(hist_real / hist_fake)) - - -def jsd(reals, fakes): - """Jensen-Shannon divergence calculation. - - Args: - reals (numpy.array): Real images. - fakes (numpy.array): Fake images. - """ - reals = reals.transpose(1, 0, 2, 3).reshape(reals.shape[1], -1) - fakes = fakes.transpose(1, 0, 2, 3).reshape(fakes.shape[1], -1) - - hist_real = np.apply_along_axis(lambda a: np.histogram(a, bins=40, range=(-1, 1))[0], 1, reals) - hist_fake = np.apply_along_axis(lambda a: np.histogram(a, bins=40, range=(-1, 1))[0], 1, fakes) - - hist_real = hist_real + 1 - hist_fake = hist_fake + 1 - - hist_real = hist_real / np.sum(hist_real) - hist_fake = hist_fake / np.sum(hist_fake) - - hist_avg = (hist_real + hist_fake) / 2 - - return 0.5 * (np.mean(np.log(hist_real / hist_avg)) + np.mean(np.log(hist_fake / hist_avg))) - - -def haar(image): - # Load the Haar cascade for face detection - face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml') - - # Convert the image to grayscale - gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) - - # Perform face detection - faces = face_cascade.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=4) - - # Draw rectangles around the detected faces - for (x, y, w, h) in faces: - cv2.rectangle(image, (x, y), (x+w, y+h), (255, 0, 0), 2) - - # Display the result - cv2.imshow('Face Detection', image) - cv2.waitKey(0) - cv2.destroyAllWindows() - - -if __name__ == '__main__': - # Load the image - image = cv2.imread('data/edface/500/03120500_000.png') - - # Perform face detection - haar(image) \ No newline at end of file -- cgit v1.3.1