Create autoencoder.py

AutoEncoder/autoencoder.py

@@ -0,0 +1,90 @@
+
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torchvision import datasets, transforms
+import matplotlib.pyplot as plt
+
+transform = transforms.Compose([transforms.ToTensor()])
+train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
+train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=64, shuffle=True)
+
+class Autoencoder(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.encoder = nn.Sequential(
+            nn.Linear(28*28, 128), nn.ReLU(),
+            nn.Linear(128, 64), nn.ReLU(),
+            nn.Linear(64, 32)
+        )
+        self.decoder = nn.Sequential(
+            nn.Linear(32, 64), nn.ReLU(),
+            nn.Linear(64, 128), nn.ReLU(),
+            nn.Linear(128, 28*28), nn.Sigmoid()
+        )
+
+    def forward(self, x):
+        x = x.view(-1, 28*28)
+        encoded = self.encoder(x)
+        decoded = self.decoder(encoded)
+        return decoded.view(-1, 1, 28, 28)
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = Autoencoder().to(device)
+criterion = nn.MSELoss()
+optimizer = optim.Adam(model.parameters(), lr=0.001)
+
+epochs = 10
+loss_history = []
+
+for epoch in range(epochs):
+    running_loss = 0
+    for data, _ in train_loader:
+        img = data.to(device)
+        output = model(img)
+        loss = criterion(output, img)
+
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+        running_loss += loss.item()
+
+    avg_loss = running_loss/len(train_loader)
+    loss_history.append(avg_loss)
+    print(f"Epoch [{epoch+1}/{epochs}], Loss: {avg_loss:.4f}")
+
+# --- VISUALIZATION SECTION ---
+
+# 1. Plot Training Loss Curve
+plt.figure(figsize=(10, 5))
+plt.plot(loss_history, marker='o', color='b')
+plt.title("Training Loss Over Epochs")
+plt.xlabel("Epoch")
+plt.ylabel("MSE Loss")
+plt.grid(True)
+plt.show()
+
+# 2. Compare Original vs Reconstructed
+model.eval()
+with torch.no_grad():
+    dataiter = iter(train_loader)
+    images, _ = next(dataiter)
+    images = images.to(device)
+    reconstructed = model(images)
+
+    images = images.cpu().numpy()
+    reconstructed = reconstructed.cpu().numpy()
+
+    fig, axes = plt.subplots(nrows=2, ncols=10, sharex=True, sharey=True, figsize=(20, 4))
+
+    for i in range(10):
+        axes[0, i].imshow(images[i].squeeze(), cmap='gray')
+        axes[0, i].set_title("Original")
+        axes[0, i].axis('off')
+
+        axes[1, i].imshow(reconstructed[i].squeeze(), cmap='gray')
+        axes[1, i].set_title("Recon")
+        axes[1, i].axis('off')
+
+    plt.tight_layout()
+    plt.show()

Variational-AutoEncoder/vae.py

@@ -0,0 +1,99 @@
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+from torchvision import datasets, transforms
+import matplotlib.pyplot as plt
+import numpy as np
+
+# Setup
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+transform = transforms.Compose([transforms.ToTensor()])
+train_loader = torch.utils.data.DataLoader(
+    datasets.MNIST('./data', train=True, download=True, transform=transform),
+    batch_size=128, shuffle=True)
+
+class VAE(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.fc1 = nn.Linear(784, 400)
+        self.mu = nn.Linear(400, 20)
+        self.logvar = nn.Linear(400, 20)
+        self.fc3 = nn.Linear(20, 400)
+        self.fc4 = nn.Linear(400, 784)
+
+    def encode(self, x):
+        h1 = F.relu(self.fc1(x))
+        return self.mu(h1), self.logvar(h1)
+
+    def reparameterize(self, mu, logvar):
+        std = torch.exp(0.5*logvar)
+        eps = torch.randn_like(std)
+        return mu + eps*std
+
+    def decode(self, z):
+        return torch.sigmoid(self.fc4(F.relu(self.fc3(z))))
+
+    def forward(self, x):
+        mu, logvar = self.encode(x.view(-1, 784))
+        z = self.reparameterize(mu, logvar)
+        return self.decode(z), mu, logvar
+
+def loss_function(recon_x, x, mu, logvar):
+    BCE = F.binary_cross_entropy(recon_x, x.view(-1, 784), reduction='sum')
+    KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
+    return BCE + KLD
+
+model = VAE().to(device)
+optimizer = optim.Adam(model.parameters(), lr=1e-3)
+
+# Training
+epochs = 15
+for epoch in range(epochs):
+    model.train()
+    train_loss = 0
+    for data, _ in train_loader:
+        data = data.to(device)
+        optimizer.zero_grad()
+        recon_batch, mu, logvar = model(data)
+        loss = loss_function(recon_batch, data, mu, logvar)
+        loss.backward()
+        train_loss += loss.item()
+        optimizer.step()
+    print(f"Epoch {epoch+1}, Avg Loss: {train_loss / len(train_loader.dataset):.4f}")
+
+# --- VISUALIZATION SECTION ---
+
+# 1. Image Generation (Sampling from Latent Space)
+model.eval()
+with torch.no_grad():
+    # Sample 10 random vectors from the normal distribution
+    sample = torch.randn(10, 20).to(device)
+    generated = model.decode(sample).cpu().view(10, 28, 28)
+
+    plt.figure(figsize=(15, 3))
+    for i in range(10):
+        plt.subplot(1, 10, i+1)
+        plt.imshow(generated[i], cmap='gray')
+        plt.axis('off')
+        plt.title("Generated")
+    plt.suptitle("New Digits Created from Random Noise")
+    plt.show()
+
+# 2. Reconstruct vs Original
+with torch.no_grad():
+    data, _ = next(iter(train_loader))
+    data = data.to(device)
+    recon, mu, logvar = model(data)
+
+    fig, axes = plt.subplots(2, 10, figsize=(15, 4))
+    for i in range(10):
+        axes[0, i].imshow(data[i].cpu().view(28, 28), cmap='gray')
+        axes[0, i].axis('off')
+        axes[1, i].imshow(recon[i].cpu().view(28, 28), cmap='gray')
+        axes[1, i].axis('off')
+    axes[0, 0].set_ylabel("Original", size=15)
+    axes[1, 0].set_ylabel("Reconstructed", size=15)
+    plt.show()
+