diff options
Diffstat (limited to 'main.py')
| -rw-r--r-- | main.py | 246 |
1 files changed, 246 insertions, 0 deletions
@@ -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') |