diff options
| -rw-r--r-- | autoencoder.py | 136 | ||||
| -rw-r--r-- | main.py | 38 |
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() @@ -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 |