testing with latent diffusion

2 files changed, 165 insertions, 9 deletions

diff --git a/autoencoder.py b/autoencoder.py
new file mode 100644
index 0000000..f16fe79
--- /dev/null
+++ b/autoencoder.py
@@ -0,0 +1,136 @@
+import torch
+from torch.utils.data import DataLoader, Dataset
+
+from torchvision import datasets, transforms
+
+import matplotlib.pyplot as plt
+
+from trainer import Trainer
+
+
+class PrintLayer(torch.nn.Module):
+    def forward(self, x):
+        print(x.shape)
+        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.latent_size = torch.prod(torch.tensor(latent_dim))
+        self.encoder = torch.nn.Sequential(
+            torch.nn.Conv2d(input_dim[0], 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.Flatten(),
+            torch.nn.Linear(64 * 7 * 7, self.latent_size),
+            torch.nn.ReLU(),
+            torch.nn.Unflatten(1, latent_dim),
+        )
+        self.decoder = torch.nn.Sequential(
+            torch.nn.Flatten(),
+            torch.nn.Linear(self.latent_size, 64 * 7 * 7),
+            torch.nn.ReLU(),
+            torch.nn.Unflatten(1, (64, 7, 7)),
+            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
+
+    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, 28, 28), latent_dim=(1, 8, 8))
+    # model.load_state_dict(torch.load('autoencoder.pth'))
+    model.to(device)
+
+    # Train model
+    trainer = Trainer()
+    lr = 1e-3
+    epochs = 10
+    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
+    criterion = torch.nn.functional.mse_loss
+    trainer.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, 4))
+        for i, j in enumerate(torch.randint(0, len(dataset), (n,))):
+            x, _ = dataset[j]
+            x = x.unsqueeze(0)
+            x_hat = model(x)
+            plt.subplot(2, n, i + 1)
+            plt.imshow(x.cpu().squeeze().numpy(), cmap='gray')
+            plt.axis('off')
+            plt.subplot(2, n, i + n + 1)
+            plt.imshow(x_hat.cpu().squeeze().numpy(), cmap='gray')
+            plt.axis('off')
+        plt.tight_layout()
+        plt.savefig('autoencoder.tmp.png')
+
+
+if __name__ == '__main__':
+    main()
diff --git a/main.py b/main.py
index 6ba7eb6..7bbbc33 100644
--- a/main.py
+++ b/main.py
@@ -12,6 +12,7 @@ import os
 import cv2
 
 from trainer import Trainer
+from autoencoder import Autoencoder
 
 
 class UNet(nn.Module):
@@ -176,6 +177,22 @@ class DiffusionDataset(Dataset):
         return len(self.dataset)
 
 
+class LatentDataset(Dataset):
+    def __init__(self, dataset, autoencoder):
+        super().__init__()
+        self.dataset = dataset
+        self.autoencoder = autoencoder
+
+    def __getitem__(self, index):
+        img, label = self.dataset[index]
+        img = img.to(self.autoencoder.device)
+        latent = self.autoencoder.encode(img.unsqueeze(0))
+        return latent, label
+
+    def __len__(self):
+        return len(self.dataset)
+
+
 def loss(y_pred, y_true):
     return nn.MSELoss()(y_pred, y_true)
 
@@ -204,12 +221,12 @@ 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.MNIST(
+    root="./data",
+    train=True,
+    download=True,
+    transform=transforms.ToTensor(),
+)
 # dataset = datasets.LFWPeople(
 #     root="./data",
 #     download=True,
@@ -219,13 +236,16 @@ ALPHA_BAR = torch.cumprod(ALPHA, dim=0)
 #     ]),
 # )
 # dataset = FolderDataset('data/lfwcrop_color/faces')
-dataset = FolderDataset('data/edface')
+# dataset = FolderDataset('data/edface')
+
+autoencoder = Autoencoder(1, 64).to(DEVICE)
+dataset = LatentDataset(dataset, autoencoder)
 
 img = dataset[0][0]
 NB_CHANNEL, IMG_SIZE, _ = img.shape
-NB_LABEL = 1
+NB_LABEL = 10
 
-EPOCHS = 100
+EPOCHS = 10
 LEARNING_RATE = 2e-4