PyTorch Autoencoder Implementation: A Step-by-Step Guide

This code implements an autoencoder in PyTorch, a popular deep learning framework. The autoencoder is trained to reconstruct input data, often used for dimensionality reduction or learning latent representations.

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

# Define the autoencoder model
class Autoencoder(nn.Module):
    def __init__(self):
        super(Autoencoder, self).__init__()
        self.encoder = nn.Sequential(
            nn.Linear(784, 128),
            nn.ReLU(),
            nn.Linear(128, 64),
            nn.ReLU(),
            nn.Linear(64, 32),
            nn.ReLU(),
        )
        self.decoder = nn.Sequential(
            nn.Linear(32, 64),
            nn.ReLU(),
            nn.Linear(64, 128),
            nn.ReLU(),
            nn.Linear(128, 784),
            nn.Sigmoid(),
        )

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

# Prepare the data
x_train = torch.tensor(x_train).float()
x_test = torch.tensor(x_test).float()
train_dataset = TensorDataset(x_train, x_train)
test_dataset = TensorDataset(x_test, x_test)
train_loader = DataLoader(train_dataset, batch_size=128, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=128, shuffle=True)

# Create the autoencoder model and optimizer
autoencoder = Autoencoder()
optimizer = Adam(autoencoder.parameters())

# Define the loss function
criterion = nn.MSELoss()

# Train the autoencoder
for epoch in range(300):
    train_loss = 0
    for batch in train_loader:
        optimizer.zero_grad()
        x, y = batch
        output = autoencoder(x)
        loss = criterion(output, y)
        loss.backward()
        optimizer.step()
        train_loss += loss.item() * x.shape[0]
    train_loss /= len(train_loader.dataset)

    test_loss = 0
    with torch.no_grad():
        for batch in test_loader:
            x, y = batch
            output = autoencoder(x)
            loss = criterion(output, y)
            test_loss += loss.item() * x.shape[0]
        test_loss /= len(test_loader.dataset)

    print(f'Epoch {epoch + 1}, Train Loss: {train_loss:.4f}, Test Loss: {test_loss:.4f}')

Explanation:

1. Model Definition: Defines the autoencoder model with encoder and decoder components using nn.Sequential for layer organization.
2. Data Preparation: Prepares training and testing datasets using TensorDataset and loads them with DataLoader for batching and shuffling.
3. Optimization: Sets up the Adam optimizer for updating model weights during training.
4. Loss Function: Defines the mean squared error (nn.MSELoss) as the loss function to measure reconstruction error.
5. Training Loop: Iterates through epochs, processes batches of data, calculates loss, performs backpropagation, and updates weights.
6. Evaluation: Tracks training and testing losses throughout epochs to monitor the autoencoder's performance.