first commit

4 files changed, 490 insertions, 0 deletions

diff --git a/.gitignore b/.gitignore
index 68bc17f..05f10ff 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,3 +1,12 @@
+# Torch datas.
+data/
+
+# Model weights.
+*.pth
+
+# Tmp pngs.
+*.tmp.png
+
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]
diff --git a/main.py b/main.py
new file mode 100644
index 0000000..7b9160a
--- /dev/null
+++ b/main.py
@@ -0,0 +1,246 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import DataLoader, Dataset
+
+from torchvision import datasets, transforms
+
+import matplotlib.pyplot as plt
+
+import numpy as np
+
+from trainer import Trainer
+
+
+class UNet(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+        # Input
+        # The input to the model is a 10 vector which represents the input image.
+        # The Input is passed through layers to generate a 1x28x28, 1x14x14, 1x7x7 tensor.
+        # -------
+        # input: 1x10
+        self.inconv1 = nn.Linear(10, 28 * 28)
+        self.inconv2 = nn.Linear(10, 14 * 14)
+        self.inconv3 = nn.Linear(10, 7 * 7)
+
+        # Encoder
+        # In the encoder, convolutional layers with the Conv2d function are used to extract features from the input image.
+        # Each block in the encoder consists of two convolutional layers followed by a max-pooling layer,
+        # with the exception of the last block which does not include a max-pooling layer.
+        # -------
+        # input: 28x28x1
+        self.e11 = nn.Conv2d(1, 64, kernel_size=3, padding=1)  # output: 28x28x64
+        self.e12 = nn.Conv2d(64, 64, kernel_size=3, padding=1)  # output: 28x28x64
+        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)  # output: 14x14x64
+
+        # input: 14x14x64
+        self.e21 = nn.Conv2d(64, 128, kernel_size=3, padding=1)  # output: 14x14x128
+        self.e22 = nn.Conv2d(128, 128, kernel_size=3, padding=1)  # output: 14x14x128
+        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)  # output: 7x7x128
+
+        # input: 7x7x128
+        self.e31 = nn.Conv2d(129, 256, kernel_size=3, padding=1)  # output: 7x7x256
+        self.e32 = nn.Conv2d(257, 256, kernel_size=3, padding=1)  # output: 7x7x256
+
+        # Decoder
+        # In the decoder, the output of the encoder is upsampled using the ConvTranspose2d function.
+        # Each block in the decoder consists of two convolutional layers followed by an upsampling layer,
+        # with the exception of the last block which does not include an upsampling layer.
+        # -------
+        # input: 7x7x256
+        self.upconv1 = nn.ConvTranspose2d(256, 128, kernel_size=2, stride=2)  # output: 14x14x128
+        self.d11 = nn.Conv2d(256, 128, kernel_size=3, padding=1)  # output: 14x14x(128x2)
+        self.d12 = nn.Conv2d(128, 128, kernel_size=3, padding=1)  # output: 14x14x128
+
+        # input: 14x14x128
+        self.upconv2 = nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2)  # output: 28x28x64
+        self.d21 = nn.Conv2d(128, 64, kernel_size=3, padding=1)  # output: 28x28x(64x2)
+        self.d22 = nn.Conv2d(64, 64, kernel_size=3, padding=1)  # output: 28x28x64
+
+        # Output
+        # The output of the decoder is passed through a convolutional layer with the Conv2d function to obtain the final output.
+        # -------
+        # input: 28x28x64
+        self.outconv = nn.Conv2d(64, 1, kernel_size=1)  # output: 28x28x1
+
+    def forward(self, x, y):
+        # Input
+        y = self.inconv3(y)
+        y = y.view(-1, 1, 7, 7)
+
+        # Encoder
+        x = F.relu(self.e11(x))
+        x1 = F.relu(self.e12(x))
+        x = self.pool1(x1)
+
+        x = F.relu(self.e21(x))
+        x2 = F.relu(self.e22(x))
+        x = self.pool2(x2)
+
+        x = torch.cat([x, y], dim=1)
+        x = F.relu(self.e31(x))
+        x = torch.cat([x, y], dim=1)
+        x = F.relu(self.e32(x))
+
+        # Decoder
+        x = self.upconv1(x)
+        x = torch.cat([x, x], dim=1)
+        x = F.relu(self.d11(x))
+        x = F.relu(self.d12(x))
+
+        x = self.upconv2(x)
+        x = torch.cat([x, x1], dim=1)
+        x = F.relu(self.d21(x))
+        x = F.relu(self.d22(x))
+
+        # Output
+        x = self.outconv(x)
+
+        return x
+
+
+DIFFU_STEPS = 10
+
+
+class MNISTDiffusionDataset(Dataset):
+    def __init__(self, train=True):
+        super().__init__()
+        self.mnist_data = datasets.MNIST(root='./data', train=train, download=True, transform=transforms.ToTensor())
+
+    def __getitem__(self, index):
+        # Get the image and the label.
+        img, label = self.mnist_data[index]
+
+        # Add noise to the image.
+        noise = np.random.normal(0, 1, (28, 28))
+        alpha = np.random.uniform(1 / DIFFU_STEPS, 1.0)
+
+        # The target is the image with the noise.
+        target = img * alpha + noise * (1 - alpha)
+
+        # The input is the image with more noise.
+        input = img * (alpha - 1 / DIFFU_STEPS) + noise * (1 - alpha + 1 / DIFFU_STEPS)
+
+        # Convert the label to a one-hot vector.
+        vector = torch.nn.functional.one_hot(torch.tensor(label), num_classes=10)
+
+        return (input.clone().detach().to(dtype=torch.float32),
+                vector.clone().detach().to(dtype=torch.float32),
+                target.clone().detach().to(dtype=torch.float32))
+
+    def __len__(self):
+        return len(self.mnist_data)
+
+
+############
+# Training #
+############
+
+# Define the device.
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+# Load the model.
+model = UNet().to(device)
+# model.load_state_dict(torch.load('model.pth'))
+
+# Define the optimizer.
+optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
+
+# Define the training dataset.
+train_dataset = MNISTDiffusionDataset(train=True)
+train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
+trainer = Trainer()
+
+# Train the model.
+trainer.train(model, train_loader, 2, optimizer, F.mse_loss)
+
+# Save the model.
+torch.save(model.state_dict(), 'model.pth')
+
+
+##############
+# Evaluation #
+##############
+
+# Load the model.
+model = UNet().to(device)
+model.load_state_dict(torch.load('model.pth'))
+
+# Set the model to evaluation mode.
+model.eval()
+
+n = 10
+fig = plt.figure(figsize=(2 * 2 * n, 3 * 2))
+gs = plt.GridSpec(nrows=3, ncols=2*2*n)
+for i in range(n):
+    # Get the i-th input and its label.
+    input, vector, label = train_dataset[i]
+
+    # Plot the input.
+    ax = fig.add_subplot(gs[0:1, 1 + 4 * i:3 + 4 * i])
+    ax.imshow(input[0], cmap='gray')
+    ax.axis('off')
+
+    # Plot the label.
+    ax = fig.add_subplot(gs[1:2, 4 * i:2 + 4 * i])
+    ax.imshow(label[0], cmap='gray')
+    ax.axis('off')
+
+    # Get the model output.
+    output = model(input.unsqueeze(0).to(device), vector.unsqueeze(0).to(device))
+
+    # Plot the model output.
+    ax = fig.add_subplot(gs[1:2, 2 + 4 * i:4 + 4 * i])
+    ax.imshow(output[0, 0].cpu().detach(), cmap='gray')
+    ax.axis('off')
+
+    # Plot the difference between the label and the model output.
+    ax = fig.add_subplot(gs[2:3, 1 + 4 * i:3 + 4 * i])
+    ax.imshow((label - output[0, 0].cpu().detach())[0], cmap='coolwarm')
+    ax.axis('off')
+
+fig.tight_layout()
+fig.savefig('diff.tmp.png')
+
+
+# Plot the evolution of the noise.
+fig = plt.figure(figsize=(n, DIFFU_STEPS))
+noises = np.random.normal(0, 1, (n, 1, 28, 28))
+noises = torch.Tensor(noises).to(device)
+vector = torch.nn.functional.one_hot(torch.tensor(range(n)), num_classes=10).to(device)
+vector = vector.clone().detach().to(dtype=torch.float32)
+
+# Apply the model multiple times.
+for i in range(DIFFU_STEPS):
+    noises = model(noises, vector)
+
+    for j in range(n):
+        ax = fig.add_subplot(DIFFU_STEPS, n, i * n + j + 1)
+        ax.imshow(noises[j, 0].cpu().detach(), cmap='gray')
+        ax.axis('off')
+
+fig.tight_layout()
+fig.savefig('diffu.tmp.png')
+
+# Plot bench of generated images.
+fig = plt.figure(figsize=(n, n))
+noises = np.random.normal(0, 1, (n * n, 1, 28, 28))
+noises = torch.Tensor(noises).to(device)
+vector = torch.nn.functional.one_hot(torch.tensor([range(n)] * n), num_classes=10).to(device)
+vector = vector.clone().detach().to(dtype=torch.float32)
+
+# Apply the model multiple times.
+for i in range(DIFFU_STEPS):
+    noises = model(noises, vector)
+
+for i in range(n * n):
+    ax = fig.add_subplot(n, n, i + 1)
+    ax.imshow(noises[i, 0].cpu().detach(), cmap='gray')
+    ax.axis('off')
+
+fig.tight_layout()
+fig.savefig('bench.tmp.png')
+
+plt.close('all')
diff --git a/requirements.txt b/requirements.txt
new file mode 100644
index 0000000..2eac55c
--- /dev/null
+++ b/requirements.txt
@@ -0,0 +1,3 @@
+torchvision==0.15.2
+matplotlib==3.8.0
+rich==13.3.5
\ No newline at end of file
diff --git a/trainer.py b/trainer.py
new file mode 100644
index 0000000..da1b3e7
--- /dev/null
+++ b/trainer.py
@@ -0,0 +1,232 @@
+import torch
+import torch.utils.data
+
+import rich.progress
+
+from typing import *
+
+
+class TrainProgress(rich.progress.Progress):
+    """A progress bar which tracks the progress of training epochs."""
+
+    def __init__(
+        self: 'TrainProgress',
+        nb_epochs: int,
+        epoch_size: int,
+        *columns: str | rich.progress.ProgressColumn,
+        console: rich.progress.Console | None = None,
+        auto_refresh: bool = True,
+        refresh_per_second: float = 10,
+        speed_estimate_period: float = 30,
+        transient: bool = False,
+        redirect_stdout: bool = True,
+        redirect_stderr: bool = True,
+        get_time: rich.progress.GetTimeCallable | None = None,
+        disable: bool = False,
+        expand: bool = False,
+    ) -> None:
+        """Initialize the progress bar.
+
+        Args:
+            nb_epochs (int): The number of epochs.
+            epoch_size (int): The size of each epoch.
+            *columns (str | rich.progress.ProgressColumn): The columns to display.
+            console (rich.progress.Console, optional): The console to use. Defaults to None.
+            auto_refresh (bool, optional): Whether to automatically refresh the progress bar. Defaults to True.
+            refresh_per_second (float, optional): The number of times to refresh the progress bar per second. Defaults to 10.
+            speed_estimate_period (float, optional): The number of seconds to use when estimating the speed. Defaults to 30.
+            transient (bool, optional): Whether to use transient mode. Defaults to False.
+            redirect_stdout (bool, optional): Whether to redirect stdout. Defaults to True.
+            redirect_stderr (bool, optional): Whether to redirect stderr. Defaults to True.
+            get_time (rich.progress.GetTimeCallable, optional): A callable which returns the current time. Defaults to None.
+            disable (bool, optional): Whether to disable the progress bar. Defaults to False.
+            expand (bool, optional): Whether to expand the progress bar. Defaults to False.
+        """
+        self.nb_epochs = nb_epochs
+        self.epoch_size = epoch_size
+        super().__init__(
+            *columns,
+            console=console,
+            auto_refresh=auto_refresh,
+            refresh_per_second=refresh_per_second,
+            speed_estimate_period=speed_estimate_period,
+            transient=transient,
+            redirect_stdout=redirect_stdout,
+            redirect_stderr=redirect_stderr,
+            get_time=get_time,
+            disable=disable,
+            expand=expand,
+        )
+        self.epoch_tasks = []
+        self.total_task = self.add_task("total", progress_type="total", total=nb_epochs*epoch_size)
+        self.values = {}
+
+    def get_renderables(self: 'TrainProgress'):
+        """Override the default renderables to display the epoch number."""
+        pad = len(f"{self.nb_epochs}")
+        for task in self.tasks:
+            # The total task.
+            if task.fields.get("progress_type") == "total":
+                self.columns = (
+                    f"Training:",
+                    rich.progress.BarColumn(),
+                    f"{len(self.epoch_tasks):{pad}}/{self.nb_epochs}",
+                    "•",
+                    rich.progress.TimeRemainingColumn(),
+                )
+
+            # The epoch tasks.
+            if task.fields.get("progress_type") == "epoch":
+                epoch_id = task.fields.get("epoch_id")
+                self.columns = (
+                    f"Epoch {epoch_id:{pad}}:",
+                    rich.progress.BarColumn(),
+                    f"{task.completed}/{task.total}",
+                    "•",
+                    rich.progress.TimeElapsedColumn(),
+                    '•',
+                    ' | '.join(f"{key}: {value[-1]:.4f} " for key, value in self.values.items()),
+                )
+
+            yield self.make_tasks_table([task])
+
+    def new_epoch(self: 'TrainProgress'):
+        """Create a new epoch task."""
+        epoch_task = self.add_task(
+            "epoch",
+            progress_type="epoch",
+            epoch_id=len(self.epoch_tasks) + 1,
+            total=self.epoch_size,
+        )
+        self.epoch_tasks.append(epoch_task)
+
+    def step(self: 'TrainProgress', count: int = 1):
+        """Advance the progress bar by the given number of steps.
+
+        Args:
+            count (int): The number of steps to advance the progress bar by.
+        """
+        if len(self.epoch_tasks) == 0 or self.tasks[self.epoch_tasks[-1]].completed == self.epoch_size:
+            self.new_epoch()
+        self.update(self.epoch_tasks[-1], advance=count)
+        self.update(self.total_task, advance=count)
+
+    def new_values(self: 'TrainProgress', **values: Any):
+        """Update the progress bar with new values.
+
+        Args:
+            **values (Any): The values to update the progress bar with.
+        """
+        for key, value in values.items():
+            current_value = self.values.get(key, [])
+            self.values[key] = current_value + [value]
+
+
+class Trainer:
+    """A class which trains models."""
+
+    def __init__(self):
+        """Initialize the trainer."""
+        self.progress: TrainProgress | None = None
+
+    def train(
+        self: 'Trainer',
+        model: torch.nn.Module,
+        train_loader: torch.utils.data.DataLoader,
+        epochs: int,
+        optimizer: torch.optim.Optimizer,
+        criterion: Callable[[torch.Tensor, torch.Tensor], torch.Tensor],
+        device: torch.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu'),
+        val_loader: torch.utils.data.DataLoader | None = None,
+    ):
+        """Train the model for the given number of epochs.
+
+        Args:
+            model (torch.nn.Module): The model to train.
+            train_loader (torch.utils.data.DataLoader): The training dataset.
+            epochs (int): The number of epochs to train the model for.
+            optimizer (torch.optim.Optimizer): The optimizer to use.
+            criterion (Callable[[torch.Tensor, torch.Tensor], torch.Tensor]): The loss function to use.
+            device (torch.device, optional): The device to use. Defaults to torch.device('cuda' if torch.cuda.is_available() else 'cpu').
+            val_loader (torch.utils.data.DataLoader, optional): The validation dataset. Defaults to None.
+        """
+        with TrainProgress(
+            nb_epochs=epochs,
+            epoch_size=len(train_loader),
+        ) as progress:
+            self.progress = progress
+
+            for _ in range(epochs):
+                self.train_epoch(
+                    model,
+                    train_loader,
+                    optimizer,
+                    criterion,
+                    device,
+                )
+                if val_loader:
+                    self.validate(
+                        model,
+                        val_loader,
+                        criterion,
+                        device,
+                    )
+
+    def train_epoch(
+        self: 'Trainer',
+        model: torch.nn.Module,
+        train_loader: torch.utils.data.DataLoader,
+        optimizer: torch.optim.Optimizer,
+        criterion: Callable[[torch.Tensor, torch.Tensor], torch.Tensor],
+        device: torch.device,
+    ):
+        """Train the model for one epoch.
+
+        Args:
+            model (torch.nn.Module): The model to train.
+            train_loader (torch.utils.data.DataLoader): The training dataset.
+            optimizer (torch.optim.Optimizer): The optimizer to use.
+            criterion (Callable[[torch.Tensor, torch.Tensor], torch.Tensor]): The loss function to use.
+            device (torch.device): The device to use.
+        """
+        model.train()
+        for batch in train_loader:
+            # Move the batch to the device.
+            batch = [b.to(device) for b in batch]
+
+            # Seprarate the inputs and labels.
+            inputs = batch[:-1]
+            labels = batch[-1]
+
+            # Train the model.
+            optimizer.zero_grad()
+            output = model(*inputs)
+            loss = criterion(output, labels)
+            loss.backward()
+            optimizer.step()
+
+            # Update the progress bar.
+            self.progress.new_values(loss=loss.item())
+            self.progress.step()
+
+    def validate(
+        self: 'Trainer',
+        model: torch.nn.Module,
+        val_loader: torch.utils.data.DataLoader,
+        criterion: Callable[[torch.Tensor, torch.Tensor], torch.Tensor],
+        device: torch.device,
+    ):
+        """Validate the model on the given validation dataset.
+
+        Args:
+            model (torch.nn.Module): The model to validate.
+            val_loader (torch.utils.data.DataLoader): The validation dataset.
+            criterion (Callable[[torch.Tensor, torch.Tensor], torch.Tensor]): The loss function to use.
+            device (torch.device): The device to use.
+        """
+        model.eval()
+        with torch.no_grad():
+            for batch in val_loader:
+                batch = [b.to(device) for b in batch]
+                output = model(batch[:-1])
+                loss = criterion(output, batch[-1])