PyTorch Autoencoder Implementation for Image Reconstruction

This code provides a detailed implementation of an Autoencoder in PyTorch for image reconstruction, using a simple linear encoder-decoder architecture. The code uses the MNIST dataset as input and aims to reconstruct the original images after encoding and decoding.

Model Architecture

The Autoencoder model consists of two main components: the encoder and the decoder. The encoder takes an input image and compresses it into a lower-dimensional representation (the code). The decoder then takes this code and attempts to reconstruct the original image.

Training Process

The model is trained using a reconstruction loss function, which measures the difference between the reconstructed image and the original image. The training process aims to minimize this loss, thereby improving the model's ability to reconstruct images.

Code

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

# Define the autoencoder model
class Autoencoder(nn.Module):
    def __init__(self):
        super(Autoencoder, self).__init__()
        self.encoder = nn.Sequential(
            nn.Linear(784, 256),
            nn.ReLU(),
            nn.Linear(256, 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, 256),
            nn.ReLU(),
            nn.Linear(256, 784),
            nn.Sigmoid()
        )

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

# Convert data to tensors and create data loaders
x_train = torch.Tensor(x_train)
x_test = torch.Tensor(x_test)
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)

# Define the autoencoder model and optimizer
autoencoder = Autoencoder()
criterion = nn.BCELoss()
optimizer = optim.Adam(autoencoder.parameters())

# Train the autoencoder
num_epochs = 300
for epoch in range(num_epochs):
    running_loss = 0.0
    for data in train_loader:
        inputs, _ = data
        optimizer.zero_grad()
        outputs = autoencoder(inputs)
        loss = criterion(outputs, inputs)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
    train_loss = running_loss / len(train_loader)
    with torch.no_grad():
        running_loss = 0.0
        for data in test_loader:
            inputs, _ = data
            outputs = autoencoder(inputs)
            loss = criterion(outputs, inputs)
            running_loss += loss.item()
        test_loss = running_loss / len(test_loader)
    print(f'Epoch {epoch+1}/{num_epochs}, Train Loss: {train_loss:.6f}, Test Loss: {test_loss:.6f}')

Explanation

1. Model Definition: The Autoencoder class defines the network structure with separate encoder and decoder modules. The encoder layers consist of linear transformations followed by ReLU activation functions. The decoder mirrors the encoder structure, ending with a sigmoid activation to produce output values between 0 and 1, suitable for image pixel values.

2. Data Preparation: The MNIST dataset is loaded and converted into PyTorch tensors. These tensors are then organized into data loaders for efficient batch training.

3. Training Loop: The code iterates through multiple epochs, processing batches of data from the training loader. For each batch, the autoencoder is used to reconstruct the images. The loss is calculated using the BCELoss function (Binary Cross-Entropy loss), which is suitable for comparing the reconstructed images to the original ones. The loss is then used to update the model's weights using the Adam optimizer.

4. Evaluation: During training, the model's performance is evaluated on a separate test set. The code prints the training and test loss for each epoch, allowing you to track the model's progress.

Key Points

• This code provides a simple but effective way to implement an Autoencoder for image reconstruction.
• The architecture can be customized by adding more layers, changing the activation functions, or using convolutional layers for better performance.
• The code showcases the key elements of training a PyTorch model, including defining the network, loss function, optimizer, and training loop.
• The code can be easily extended to other image datasets or used for various tasks, such as dimensionality reduction or image denoising.