PyTorch Autoencoder

An autoencoder (Autoencoder, AE)is a neural network for unsupervised learning that compresses input data into a low-dimensional latent space through learning, and then reconstructs the original data from the compressed representation.

Autoencoders are widely used in scenarios such as data dimensionality reduction, feature extraction, anomaly detection, image denoising, and generative models.

1. Basic Principles of Autoencoder

The basic structure of an autoencoder consists of three parts:

Encoder: maps input data \(x\) to a low-dimensional latent representation \(z\)
Latent Space: the low-dimensional vector output by the encoder, also called the bottleneck layer
Decoder: reconstructs the latent representation \(z\) into the output \(\hat{x}\)

1.1 Network Structure

The goal of an autoencoder is to make the output \(\hat{x}\) as close as possible to the input \(x\):

\[ \min_{\theta, \phi} \frac{1}{n} \sum_{i=1}^{n} \| x_i - D_\phi(E_\theta(x_i)) \|^2 \]

where \(\theta\) are the encoder parameters and \(\phi\) are the decoder parameters.

1.2 Dimensionality Reduction Effect

Autoencoders learn a compressed representation of data by forcing the data through a bottleneck layer smaller than the input dimension. This compression preserves the main information of the data.

Compared with Principal Component Analysis (PCA), autoencoders can learn nonlinear dimensionality reduction and capture more complex data structures.

2. Basic Autoencoder Implementation

2.1 Simple Autoencoder

Example

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset

# ── Autoencoder model ─────────────────────────────────
class Autoencoder(nn.Module):
"""
Basic autoencoder: symmetric structure
"""
def __init__(self, input_dim, hidden_dim, latent_dim):
super().__init__()

# Encoder
self.encoder = nn.Sequential(
nn.Linear(input_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, latent_dim), # Bottleneck layer
)

# Decoder
self.decoder = nn.Sequential(
nn.Linear(latent_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, input_dim),
)

def forward(self, x):
z = self.encoder(x)
x_recon = self.decoder(z)
return x_recon

def encode(self, x):
"""Encoding: obtain latent representation"""
return self.encoder(x)

def decode(self, z):
"""Decoding: reconstruct from latent representation"""
return self.decoder(z)

# ── Usage example ──────────────────────────────────────
INPUT_DIM = 784 # For example, after flattening an MNIST image
HIDDEN_DIM = 256
LATENT_DIM = 32 # Latent space dimension, much smaller than input dimension

model = Autoencoder(INPUT_DIM, HIDDEN_DIM, LATENT_DIM)
print(f"Input dimension: {INPUT_DIM}")
print(f"Latent dimension: {LATENT_DIM}")
print(f"Compression ratio: {INPUT_DIM / LATENT_DIM:.1f}x")

# View parameter count
total_params = sum(p.numel() for p in model.parameters())
print(f"Total Parameters: {total_params:,}")

2.2 Convolutional Autoencoder

For image data, autoencoders using convolutional layers perform better:

Example

import torch
import torch.nn as nn

class ConvAutoencoder(nn.Module):
"""
Convolutional autoencoder: suitable for images
"""
def __init__(self, channels=3, latent_dim=128):
super().__init__()

# Encoder: gradually reduce size, increase number of channels
# Input: (batch, channels, 64, 64)
self.encoder = nn.Sequential(
# 32 -> 16
nn.Conv2d(channels, 32, kernel_size=3, stride=2, padding=1),
nn.ReLU(),
# 16 -> 8
nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1),
nn.ReLU(),
# 8 -> 4
nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1),
nn.ReLU(),
# 4 -> 2
nn.Conv2d(128, 256, kernel_size=3, stride=2, padding=1),
nn.ReLU(),
)

# Latent space mapping
self.to_latent = nn.AdaptiveAvgPool2d((1, 1))

# Decoder: gradually increase size
# Input: (batch, 256, 2, 2)
self.from_latent = nn.Sequential(
nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True),
nn.Conv2d(256, 128, kernel_size=3, padding=1),
nn.ReLU(),
nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True),
nn.Conv2d(128, 64, kernel_size=3, padding=1),
nn.ReLU(),
nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True),
nn.Conv2d(64, 32, kernel_size=3, padding=1),
nn.ReLU(),
nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True),
nn.Conv2d(32, channels, kernel_size=3, padding=1),
nn.Sigmoid() # Output [0, 1]
)

def forward(self, x):
z = self.encode(x)
x_recon = self.decode(z)
return x_recon

def encode(self, x):
"""Encode"""
features = self.encoder(x)
z = self.to_latent(features)
z = z.view(z.size(0), -1) # (batch, 256)
return z

def decode(self, z):
"""Decode"""
# Reshape vector to feature map
batch_size = z.size(0)
z = z.view(batch_size, 256, 1, 1)
z = z.expand(-1, -1, 2, 2) # Upsample to 2x2
x_recon = self.from_latent(z)
return x_recon

# Test
model = ConvAutoencoder(channels=3, latent_dim=128)
x = torch.randn(4, 3, 64, 64)
x_recon = model(x)

print(f"Input shape: {x.shape}")
print(f"Output shape: {x_recon.shape}")
print(f"Latent vector shape: {model.encode(x).shape}")

2.3 Training and Reconstruction

Example

import torch
import torch.nn as nn
import torch.optim as optim

# ── Training configuration ─────────────────────────────────────
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = ConvAutoencoder(channels=3, latent_dim=128).to(device)
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=1e-3)

# ── Training loop ─────────────────────────────────────
def train_autoencoder(model, dataloader, criterion, optimizer, num_epochs=10):
model.train()
for epoch in range(num_epochs):
total_loss = 0
for batch in dataloader:
images = batch[0].to(device)

# Forward propagation
outputs = model(images)
loss = criterion(outputs, images)

# Backpropagation
optimizer.zero_grad()
loss.backward()
optimizer.step()

total_loss += loss.item()

avg_loss = total_loss / len(dataloader)
print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {avg_loss:.6f}")

return model

# Assume a data loader already exists
# train_autoencoder(model, train_loader, criterion, optimizer, num_epochs=10)
print("Autoencoder training complete!")

3. Denoising Autoencoder (DAE)

Denoising Autoencoder (DAE)adds noise to the input during training, then learns to remove the noise and recover the original input. This enables the model to learn more robust feature representations.

3.1 Denoising Autoencoder Implementation

Example

import torch
import torch.nn as nn

class DenoisingAutoencoder(nn.Module):
"""
Denoising autoencoder
"""
def __init__(self, input_dim, hidden_dim, latent_dim):
super().__init__()

self.encoder = nn.Sequential(
nn.Linear(input_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, latent_dim),
)

self.decoder = nn.Sequential(
nn.Linear(latent_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, input_dim),
nn.Sigmoid() # Output [0, 1]
)

def forward(self, x):
z = self.encode(x)
return self.decode(z)

def encode(self, x):
return self.encoder(x)

def decode(self, z):
return self.decoder(z)

def add_noise(x, noise_factor=0.3):
"""
Add Gaussian noise
"""
noise = torch.randn_like(x) * noise_factor
noisy_x = x + noise
return torch.clamp(noisy_x, 0.0, 1.0)

# Train denoising autoencoder
def train_dae(model, dataloader, noise_factor=0.3, lr=1e-3):
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=lr)

model.train()
for epoch in range(10):
for batch in dataloader:
images = batch[0]

# Add noise
noisy_images = add_noise(images, noise_factor)
noisy_images = noisy_images.to(next(model.parameters()).device)
images = images.to(next(model.parameters()).device)

# Forward propagation
outputs = model(noisy_images)
loss = criterion(outputs, images) # Compare with original image, not the noisy image

# Backpropagation
optimizer.zero_grad()
loss.backward()
optimizer.step()

return model

3.2 Other Noise Types

Example

def salt_pepper_noise(x, prob=0.1):
"""Salt-and-pepper noise"""
random_mask = torch.rand_like(x)
noisy = x.clone()
noisy[random_mask < prob / 2] = 0.0
noisy[random_mask > 1 - prob / 2] = 1.0
return noisy

def mask_noise(x, prob=0.1):
"""Masking noise (randomly set to zero)"""
mask = torch.rand_like(x) > prob
return x * mask.float()

def dropout_noise(x, rate=0.2):
"""Dropout noise"""
mask = torch.rand_like(x) > rate
return x * mask.float() / (1 - rate)

4. Variational Autoencoder (VAE)

Variational Autoencoder (VAE)is a generative model that encodes data into a probability distribution in the latent space rather than a fixed vector. This allows us to sample from the latent space to generate new data.

4.1 Core Principles of VAE

The key innovation of VAE is learning the probability distribution of latent variables:

The encoder outputs mean \(\mu\) and standard deviation \(\sigma\)
Sample the latent vector \(z\) from the normal distribution \(\mathcal{N}(\mu, \sigma)\)
The decoder reconstructs data from \(z\)

To achieve a differentiable sampling process, theReparameterization Trick is used.：

\[ z = \mu + \sigma \cdot \epsilon, \quad \epsilon \sim \mathcal{N}(0, 1) \]

4.2 VAE Implementation

Example

import torch
import torch.nn as nn
import torch.optim as optim

class VAE(nn.Module):
"""
Variational autoencoder
"""
def __init__(self, input_dim, hidden_dim, latent_dim):
super().__init__()

# Encoder: output mean and variance
self.encoder = nn.Sequential(
nn.Linear(input_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, hidden_dim),
nn.ReLU(),
)
self.fc_mu = nn.Linear(hidden_dim, latent_dim)
self.fc_logvar = nn.Linear(hidden_dim, latent_dim)

# Decoder
self.decoder = nn.Sequential(
nn.Linear(latent_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, input_dim),
nn.Sigmoid()
)

def encode(self, x):
h = self.encoder(x)
mu = self.fc_mu(h)
logvar = self.fc_logvar(h)
return mu, logvar

def reparameterize(self, mu, logvar):
"""Reparameterization trick"""
std = torch.exp(0.5 * logvar)
eps = torch.randn_like(std)
return mu + eps * std

def decode(self, z):
return self.decoder(z)

def forward(self, x):
mu, logvar = self.encode(x)
z = self.reparameterize(mu, logvar)
x_recon = self.decode(z)
return x_recon, mu, logvar

def vae_loss(x_recon, x, mu, logvar, beta=1.0):
"""
VAE loss function
Reconstruction loss + KL divergence
"""
# Reconstruction loss
recon_loss = nn.functional.mse_loss(x_recon, x, reduction='sum')

# KL divergence: -0.5 * sum(1 + log(sigma^2) - mu^2 - sigma^2)
kl_loss = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())

return recon_loss + beta * kl_loss, recon_loss, kl_loss

# Usage example
INPUT_DIM = 784
HIDDEN_DIM = 256
LATENT_DIM = 2 # Two-dimensional latent space for easy visualization

model = VAE(INPUT_DIM, HIDDEN_DIM, LATENT_DIM)

# Test
x = torch.randn(32, 784)
x_recon, mu, logvar = model(x)

print(f"Input shape: {x.shape}")
print(f"Reconstructed shape: {x_recon.shape}")
print(f"Mean shape: {mu.shape}") # (32, 2)
print(f"Variance shape: {logvar.shape}") # (32, 2)

4.3 VAE Generation and Visualization

Example

import matplotlib.pyplot as plt

def visualize_latent_space(model, dataloader, device):
"""Visualizing the Latent Space"""
model.eval()
all_mu = []
all_labels = []

with torch.no_grad():
for batch in dataloader:
images, labels = batch[0].to(device), batch[1]
mu, _ = model.encode(images)
all_mu.append(mu.cpu())
all_labels.append(labels)

all_mu = torch.cat(all_mu, dim=0).numpy()
all_labels = torch.cat(all_labels, dim=0).numpy()

plt.figure(figsize=(10, 8))
scatter = plt.scatter(all_mu[:, 0], all_mu[:, 1], c=all_labels,
cmap='tab10', alpha=0.5, s=10)
plt.colorbar(scatter)
plt.xlabel('Latent Dimension 1')
plt.ylabel('Latent Dimension 2')
plt.title('VAE Latent Space')
plt.show()

def generate_from_latent(model, z, device):
"""Generating Images from Latent Vectors"""
model.eval()
with torch.no_grad():
z = z.to(device)
generated = model.decode(z)
return generated

def interpolate_latent(model, z1, z2, steps=10, device):
"""Latent Space Interpolation Generation"""
model.eval()

# Linear interpolation
alphas = torch.linspace(0, 1, steps)
interpolated = []

with torch.no_grad():
for alpha in alphas:
z = z1 * (1 - alpha) + z2 * alpha
generated = model.decode(z)
interpolated.append(generated)

return torch.cat(interpolated, dim=0)

# Generate new images
def generate_new_images(model, num_images, latent_dim, device):
"""Generate new images from random latent vectors"""
model.eval()
with torch.no_grad():
# Sample from standard normal distribution
z = torch.randn(num_images, latent_dim).to(device)
generated = model.decode(z)
return generated

The latent space of VAE is continuous, allowing interpolation in the latent space to generate smoothly transitioning images. However, images generated by VAE are often blurry because it optimizes a lower bound rather than the exact log-likelihood.

5. Sparse Autoencoder

Sparse AutoencoderAdd a sparsity constraint to the loss function, limiting the number of activations in the latent vector. This enables the model to learn more meaningful features.

5.1 Sparse Autoencoder Implementation

< h2 class="example">Example

import torch
import torch.nn as nn
import torch.nn.functional as F

class SparseAutoencoder(nn.Module):
"""
Sparse Autoencoder
"""
def __init__(self, input_dim, hidden_dim, latent_dim):
super().__init__()

# Encoder
self.encoder = nn.Sequential(
nn.Linear(input_dim, hidden_dim),
nn.ReLU(),
)
# Latent layer
self.bottleneck = nn.Linear(hidden_dim, latent_dim)

# Decoder
self.decoder = nn.Sequential(
nn.Linear(latent_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, input_dim),
nn.Sigmoid()
)

def forward(self, x):
h = self.encoder(x)
z = self.bottleneck(h)
z_activated = F.relu(z) # Sparse activation
x_recon = self.decoder(z_activated)
return x_recon, z_activated

def sparse_loss(z, rho=0.05, beta=1.0):
"""
Sparse loss: KL divergence
rho: target sparsity (e.g., 0.05 means only 5% of neurons should be active)
beta: sparsity term weight
"""
# Compute average activation
rho_hat = torch.mean(z, dim=0)

# KL divergence
kl = rho * torch.log(rho / (rho_hat + 1e-8)) + \
(1 - rho) * torch.log((1 - rho) / (1 - rho_hat + 1e-8))

return beta * torch.sum(kl)

def total_sparse_loss(x_recon, x, z, rho=0.05, beta=1.0):
"""Total loss = reconstruction loss + sparsity loss"""
recon_loss = F.mse_loss(x_recon, x)
sparsity = sparse_loss(z, rho, beta)
return recon_loss + sparsity

6. Sequence-to-Sequence Autoencoder

For sequence data (such as text, time series), use RNN/LSTM as the encoder and decoder.

6.1 Sequence Autoencoder Implementation

Example

import torch
import torch.nn as nn

class Seq2SeqAutoencoder(nn.Module):
"""
Sequence-to-sequence autoencoder: used for sequence data
"""
def __init__(self, input_size, hidden_size, latent_size, num_layers=2):
super().__init__()

# Encoder LSTM
self.encoder = nn.LSTM(
input_size=input_size,
hidden_size=hidden_size,
num_layers=num_layers,
batch_first=True,
bidirectional=True
)

# Latent space mapping
# Bidirectional LSTM output is hidden_size * 2
self.to_latent = nn.Linear(hidden_size * 2, latent_size)
self.from_latent = nn.Linear(latent_size, hidden_size * 2)

# Decoder LSTM
self.decoder = nn.LSTM(
input_size=hidden_size,
hidden_size=hidden_size,
num_layers=num_layers,
batch_first=True,
bidirectional=True
)

# Output mapping
self.output_proj = nn.Linear(hidden_size * 2, input_size)

def forward(self, x):
# Encode
_, (h_n, _) = self.encoder(x)

# Concatenate bidirectional final hidden states
# h_n: (num_layers * 2, batch, hidden_size)
h_forward = h_n[-2]
h_backward = h_n[-1]
h_combined = torch.cat([h_forward, h_backward], dim=-1)

# Map to latent space
z = self.to_latent(h_combined)

# Map back from latent space
decoder_init = self.from_latent(z)
decoder_init = decoder_init.view(2, decoder_init.size(0), -1) # (2, batch, hidden_size*2)

# Decode (using original input length)
decoder_output, _ = self.decoder(
x,
(decoder_init, torch.zeros_like(decoder_init))
)

# Output mapping
output = self.output_proj(decoder_output)
return output, z

# Usage example
model = Seq2SeqAutoencoder(
input_size=128, # Input feature dimension
hidden_size=256, # LSTM hidden dimension
latent_size=64, # Latent space dimension
num_layers=2
)

# Test
x = torch.randn(8, 20, 128) # (batch, seq_len, input_size)
output, z = model(x)

print(f"Input shape: {x.shape}") # (8, 20, 128)
print(f"Output shape: {output.shape}") # (8, 20, 128)
print(f"Latent vector shape: {z.shape}") # (8, 64)

7. Application Scenarios of Autoencoders

7.1 Anomaly Detection

Autoencoders can be used to detect anomalous data. Normal data has small reconstruction error, while anomalous data has large reconstruction error:

Example

import torch
import torch.nn as nn

def detect_anomalies(model, data_loader, threshold=None, device='cpu'):
"""
Detecting anomalies using autoencoders
"""
model.eval()
reconstruction_errors = []

with torch.no_grad():
for batch in data_loader:
images = batch[0].to(device)
outputs = model(images)

# Compute reconstruction error (mean squared error)
errors = torch.mean((outputs - images) ** 2, dim=(1, 2, 3))
reconstruction_errors.extend(errors.cpu().numpy())

reconstruction_errors = torch.tensor(reconstruction_errors)

# If no threshold is given, use statistical method
if threshold is None:
# Use 95% percentile
threshold = torch.quantile(reconstruction_errors, 0.95).item()

# Mark anomalies
anomalies = reconstruction_errors > threshold

return anomalies, reconstruction_errors, threshold

# Train anomaly detection model
def train_anomaly_detector(normal_data_loader):
"""Train the autoencoder using only normal data"""
model = ConvAutoencoder(channels=1, latent_dim=32)
criterion = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

model.train()
for epoch in range(10):
for batch in normal_data_loader:
images = batch[0]
outputs = model(images)
loss = criterion(outputs, images)

optimizer.zero_grad()
loss.backward()
optimizer.step()

return model

7.2 Image Colorization and Style Transfer

Example

class ColorizationAutoencoder(nn.Module):
"""
Image colorization autoencoder
Input: grayscale image (batch, 1, H, W)
Output: color image (batch, 2, H, W) (ab color space)
"""
def __init__(self):
super().__init__()

# Encoder: progressively extract features
self.encoder = nn.Sequential(
nn.Conv2d(1, 64, kernel_size=4, stride=2, padding=1), # H/2
nn.ReLU(),
nn.Conv2d(64, 128, kernel_size=4, stride=2, padding=1), # H/4
nn.BatchNorm2d(128),
nn.ReLU(),
nn.Conv2d(128, 256, kernel_size=4, stride=2, padding=1), # H/8
nn.BatchNorm2d(256),
nn.ReLU(),
)

# Decoder: upsample to generate colors
self.decoder = nn.Sequential(
nn.ConvTranspose2d(256, 128, kernel_size=4, stride=2, padding=1),
nn.BatchNorm2d(128),
nn.ReLU(),
nn.ConvTranspose2d(128, 64, kernel_size=4, stride=2, padding=1),
nn.BatchNorm2d(64),
nn.ReLU(),
nn.ConvTranspose2d(64, 32, kernel_size=4, stride=2, padding=1),
nn.ReLU(),
nn.Conv2d(32, 2, kernel_size=3, padding=1), # Output ab channels
nn.Sigmoid() # ab channel range [0, 1]
)

def forward(self, x):
z = self.encoder(x)
color = self.decoder(z)
return color

7.3 Data Dimensionality Reduction Visualization

Example

import matplotlib.pyplot as plt

def visualize_latent_2d(model, dataloader, device, num_samples=1000):
"""
Use an autoencoder to reduce data to 2 dimensions for visualization
"""
model.eval()
all_latents = []
all_labels = []

with torch.no_grad():
count = 0
for batch in dataloader:
if count >= num_samples:
break
images, labels = batch[0], batch[1]
images = images.to(device)

# If it's a 2D AE, use directly
# If the dimensionality is higher, project to 2D first
z = model.encode(images)

all_latents.append(z.cpu())
all_labels.append(labels)
count += images.size(0)

latents = torch.cat(all_latents, dim=0)[:num_samples].numpy()
labels = torch.cat(all_labels, dim=0)[:num_samples].numpy()

plt.figure(figsize=(10, 8))
scatter = plt.scatter(latents[:, 0], latents[:, 1], c=labels,
cmap='tab10', alpha=0.6, s=20)
plt.colorbar(scatter)
plt.xlabel('Latent Dimension 1')
plt.ylabel('Latent Dimension 2')
plt.title('Autoencoder 2D Latent Space Visualization')
plt.show()

8. API Quick Reference

8.1 Common Autoencoder Types

Type	Characteristics	Applicable Scenarios
Basic Autoencoder	Simple symmetric structure	Dimensionality reduction, feature extraction
Convolutional Autoencoder	Uses convolutional layers to preserve spatial structure	Image processing
Denoising Autoencoder	Learns to remove noise	Image denoising, robust features
Variational Autoencoder	Learns probability distributions, can generate new data	Generative models, data generation
Sparse Autoencoder	Sparsity constraint, learns interpretable features	Feature disentanglement, interpretability
Sequence Autoencoder	Uses RNN/LSTM to process sequences	Text, time series

8.2 Loss Function Selection

Task	Recommended loss function
Image reconstruction	MSELoss、SSIMLoss
Binary images	BCELoss、BCEWithLogitsLoss
Text reconstruction	CrossEntropyLoss
VAE	MSE + KL Divergence
Anomaly detection	MSE、MAE

8.3 Latent Dimension Selection

数据维度低（<100维）
    -> 潜在维度设为 2~10

数据维度中等（100~1000维）
    -> 潜在维度设为 10~50

数据维度高（>1000维）
    -> 潜在维度设为 50~200

生成任务（VAE）
    -> 潜在维度 2~32（便于采样和可视化）

异常检测
    -> 潜在维度 16~64（保留足够信息检测异常）

Other extensions