Autoencoder architecture in PyTorch

Introduction

An autoencoder learns to copy its input to its output by compressing data into a smaller form and then reconstructing it. This helps find important features in data without labels.

To reduce the size of images or data for easier storage or faster processing.

To remove noise from pictures or signals by learning clean versions.

To learn useful features from data when you don't have labels.

To generate new data similar to the input by sampling from the compressed form.

Syntax

PyTorch

class Autoencoder(nn.Module):
    def __init__(self, input_size, hidden_size, code_size):
        super().__init__()
        self.encoder = nn.Sequential(
            nn.Linear(input_size, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, code_size),
            nn.ReLU()
        )
        self.decoder = nn.Sequential(
            nn.Linear(code_size, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, input_size),
            nn.Sigmoid()
        )
    def forward(self, x):
        code = self.encoder(x)
        reconstructed = self.decoder(code)
        return reconstructed

The encoder compresses the input into a smaller code.

The decoder tries to rebuild the original input from the code.

Examples

A very simple autoencoder for 28x28 images flattened to 784 features.

PyTorch

class SimpleAutoencoder(nn.Module):
    def __init__(self):
        super().__init__()
        self.encoder = nn.Linear(784, 32)
        self.decoder = nn.Linear(32, 784)
    def forward(self, x):
        code = torch.relu(self.encoder(x))
        reconstructed = torch.sigmoid(self.decoder(code))
        return reconstructed

A deeper autoencoder with multiple layers compressing to 3 features.

PyTorch

class DeepAutoencoder(nn.Module):
    def __init__(self):
        super().__init__()
        self.encoder = nn.Sequential(
            nn.Linear(784, 128),
            nn.ReLU(),
            nn.Linear(128, 64),
            nn.ReLU(),
            nn.Linear(64, 12),
            nn.ReLU(),
            nn.Linear(12, 3)
        )
        self.decoder = nn.Sequential(
            nn.Linear(3, 12),
            nn.ReLU(),
            nn.Linear(12, 64),
            nn.ReLU(),
            nn.Linear(64, 128),
            nn.ReLU(),
            nn.Linear(128, 784),
            nn.Sigmoid()
        )
    def forward(self, x):
        code = self.encoder(x)
        reconstructed = self.decoder(code)
        return reconstructed

Sample Model

This program trains a simple autoencoder on MNIST digits for one batch and prints the loss. The loss shows how well the model reconstructs the input.

PyTorch

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader

# Define Autoencoder
class Autoencoder(nn.Module):
    def __init__(self):
        super().__init__()
        self.encoder = nn.Sequential(
            nn.Linear(28*28, 64),
            nn.ReLU(),
            nn.Linear(64, 12),
            nn.ReLU(),
            nn.Linear(12, 3)
        )
        self.decoder = nn.Sequential(
            nn.Linear(3, 12),
            nn.ReLU(),
            nn.Linear(12, 64),
            nn.ReLU(),
            nn.Linear(64, 28*28),
            nn.Sigmoid()
        )
    def forward(self, x):
        code = self.encoder(x)
        reconstructed = self.decoder(code)
        return reconstructed

# Load MNIST dataset
transform = transforms.Compose([transforms.ToTensor(), transforms.Lambda(lambda x: x.view(-1))])
mnist_train = datasets.MNIST(root='.', train=True, download=True, transform=transform)
train_loader = DataLoader(mnist_train, batch_size=128, shuffle=True)

# Initialize model, loss, optimizer
model = Autoencoder()
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# Train for 1 epoch
model.train()
for batch_idx, (data, _) in enumerate(train_loader):
    optimizer.zero_grad()
    output = model(data)
    loss = criterion(output, data)
    loss.backward()
    optimizer.step()
    if batch_idx == 0:
        print(f'Batch {batch_idx} Loss: {loss.item():.4f}')
        break

Important Notes

Use a sigmoid activation at the output if input data is normalized between 0 and 1.

The size of the code (bottleneck) controls how much the data is compressed.

Training longer and with more data improves reconstruction quality.

Summary

An autoencoder learns to compress and then reconstruct data.

It has two parts: encoder (compress) and decoder (rebuild).

It helps find important features without needing labels.