From 3b82ce5f658ab27ed6e6eaddf239c553f9431b3e Mon Sep 17 00:00:00 2001 From: gdamms Date: Thu, 6 Jun 2024 10:54:20 +0200 Subject: working ldm --- autoencoder.py | 52 ++++++++++++++++++++++++++++++++++------------------ main.py | 52 ++++++++++++++++++++++++---------------------------- 2 files changed, 58 insertions(+), 46 deletions(-) diff --git a/autoencoder.py b/autoencoder.py index 645f628..58326fe 100644 --- a/autoencoder.py +++ b/autoencoder.py @@ -21,9 +21,8 @@ class Autoencoder(torch.nn.Module): 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.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(), @@ -37,16 +36,19 @@ class Autoencoder(torch.nn.Module): 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.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(), - torch.nn.Unflatten(1, latent_dim), ) self.decoder = torch.nn.Sequential( - torch.nn.Flatten(), - torch.nn.Linear(self.latent_size, 64 * 7 * 7), + # 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.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), @@ -57,7 +59,7 @@ class Autoencoder(torch.nn.Module): torch.nn.ReLU(), torch.nn.Conv2d(16, 16, kernel_size=3, padding=1), torch.nn.ReLU(), - torch.nn.Conv2d(16, input_dim[0], kernel_size=3, padding=1), + torch.nn.Conv2d(16, input_dim, kernel_size=3, padding=1), torch.nn.Sigmoid(), ) @@ -102,14 +104,14 @@ def main(): 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 = Autoencoder(input_dim=1, latent_dim=1) + model.load_state_dict(torch.load('mnist_autoencoder.pth')) model.to(device) # Train model trainer = Trainer() lr = 1e-3 - epochs = 10 + epochs = 1 optimizer = torch.optim.Adam(model.parameters(), lr=lr) criterion = torch.nn.functional.binary_cross_entropy trainer.train(model, dataloader, epochs, optimizer, criterion) @@ -120,17 +122,31 @@ def main(): # Visualize results n = 10 with torch.no_grad(): - plt.figure(figsize=(2*n, 4)) + 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_hat = model(x) - plt.subplot(2, n, i + 1) - plt.imshow(x.cpu().squeeze().numpy(), cmap='gray') + 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(2, n, i + n + 1) - plt.imshow(x_hat.cpu().squeeze().numpy(), cmap='gray') + 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') diff --git a/main.py b/main.py index 444292f..72ea15e 100644 --- a/main.py +++ b/main.py @@ -142,15 +142,20 @@ def p_xt_1_xt(model, xt, t, vec): class DiffusionDataset(Dataset): - def __init__(self, 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 @@ -177,22 +182,6 @@ 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(DEVICE) - latent = self.autoencoder.encode(img.unsqueeze(0)).squeeze(0) - return latent, label - - def __len__(self): - return len(self.dataset) - - def loss(y_pred, y_true): return nn.MSELoss()(y_pred, y_true) @@ -238,16 +227,18 @@ dataset = datasets.MNIST( # dataset = FolderDataset('data/lfwcrop_color/faces') # dataset = FolderDataset('data/edface') -autoencoder = Autoencoder((1, 28, 28), (1, 8, 8)).to(DEVICE) +autoencoder = None +autoencoder = Autoencoder(1, 1).to(DEVICE) autoencoder.load_state_dict(torch.load('autoencoder.pth')) autoencoder.eval() -dataset = LatentDataset(dataset, autoencoder) -img = dataset[0][0] +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 = 0 +EPOCHS = 1 LEARNING_RATE = 2e-4 @@ -262,7 +253,7 @@ if __name__ == '__main__': # Load the model. model = UNet().to(DEVICE) try: - model.load_state_dict(torch.load('model.pth')) + model.load_state_dict(torch.load('mnist_latent_model.pth')) except FileNotFoundError: print("No model found, training a new one.") pass @@ -271,7 +262,7 @@ if __name__ == '__main__': optimizer = torch.optim.Adam(model.parameters(), lr=LEARNING_RATE) # Define the training dataset. - train_dataset = DiffusionDataset(dataset) + train_dataset = DiffusionDataset(dataset, autoencoder) train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True, num_workers=4, persistent_workers=True) trainer = Trainer() @@ -291,7 +282,10 @@ if __name__ == '__main__': with torch.no_grad(): # Forward diffusion img, label = dataset[np.random.randint(0, len(dataset))] - img = img.to(DEVICE) * 2 - 1 + 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)] @@ -322,7 +316,6 @@ if __name__ == '__main__': plt.suptitle("Forward diffusion") plt.tight_layout() plt.savefig("forward_diffusion.tmp.png") - exit() # Backward diffusion @@ -345,7 +338,7 @@ if __name__ == '__main__': 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(dataset.autoencoder.decode(x[class_i].unsqueeze(0)).squeeze(0))) + plt.imshow(tensor_to_image(x[class_i])) plt.axis("off") plt.tight_layout() plt.savefig("backward_diffusion.tmp.png") @@ -368,10 +361,13 @@ if __name__ == '__main__': 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 = train_dataset.autoencoder.decode(img.unsqueeze(0)).squeeze(0) plt.subplot(n_classes, nb_plots, id + 1) - plt.imshow(tensor_to_image(dataset.autoencoder.decode(x[id].unsqueeze(0)).squeeze(0))) + plt.imshow(tensor_to_image(img)) plt.axis("off") plt.tight_layout() plt.savefig("benchmark.tmp.png") - # plt.show() \ No newline at end of file + # plt.show() -- cgit v1.3.1