Build an Image Autoencoder with PyTorch

Train a neural network to compress and reconstruct handwritten digits.

Introduction

30 Second Summary

Think about how you compress a photo on your phone to send it in a message. The image gets smaller but still looks recognizable. Neural networks can learn to do this same kind of compression automatically.

In this project, you will build an autoencoder neural network that learns to compress handwritten digit images into a tiny representation and reconstruct them back.

What You'll Build

A fully-connected autoencoder built with PyTorch that compresses and reconstructs handwritten digit images from the MNIST dataset.

By the end of this project, you'll have:

  • A trained autoencoder that compresses 784-pixel images down to just 32 numbers and reconstructs them.
  • Side-by-side visualizations comparing original and reconstructed digit images to prove the network learned meaningful compression.
  • A latent space scatter plot revealing how the encoder organizes different digits into clusters.
  • Secret Mission: Transform your autoencoder into a denoising autoencoder that removes noise from corrupted images.

Are there any prerequisites?

You should have basic Python knowledge and some familiarity with neural network concepts like layers, training, and loss functions. No prior PyTorch experience is needed.

Set Up Your Environment

Before you can build an autoencoder, you need the right tools installed. PyTorch is the deep learning framework that provides the building blocks for your neural network.

In this step, you will install PyTorch and matplotlib, create your project file, and set up the configuration constants that control how your autoencoder trains.

In this step, get ready to:
  • Install PyTorch, torchvision, and matplotlib.
  • Create a project folder and autoencoder.py file.
  • Add imports and configuration constants.
Install PyTorch and matplotlib

PyTorch requires Python 3.10 or later. Since Python is already on your machine, you can install the deep learning libraries directly.

  • Press Cmd+Space (macOS) or the Windows key (Windows) to open your search bar.
  • Type Terminal and press Enter to open it.
  • Install PyTorch and torchvision by running this command:
pip3 install torch torchvision --index-url https://download.pytorch.org/whl/cpu

Why the CPU version?

This command installs the CPU-only build of PyTorch. It is smaller and faster to download than the GPU version.

For this project, CPU is more than enough to train on the MNIST dataset in a few minutes.

  • Install matplotlib by running:
pip install matplotlib
  • Verify all packages installed correctly. Open a Python prompt by running:
python3
  • In the Python prompt (you will see >>>), type these lines one at a time and press Enter after each:
import torch
print(torch.__version__)
import torchvision
import matplotlib

✔️ I see a version number printed

You should see something like 2.12.0 printed after the print(torch.__version__) line, and no errors from the other imports.

  • Type exit() to leave the Python prompt.

ⓧ I see a ModuleNotFoundError

This means one of the packages did not install correctly.

  • Type exit() to leave the Python prompt.
  • Check which package failed (the error message names it) and re-run the corresponding pip install command.
  • If torch failed, make sure you copied the full command including --index-url.
  • Confirm your Python version is 3.10 or later by running python3 --version in your terminal.

Install taking too long or failing?

The PyTorch CPU package is around 200 MB. On slower connections, this can take a few minutes.

If you see a permissions error, try adding --user to the end of the pip command. If pip3 is not found, try pip instead.

help me troubleshoot my PyTorch installation

Create your project folder and file
  • Move to your Desktop by running this command:
cd ~/Desktop
  • Create your project folder and move into it by running these commands:
mkdir autoencoder-project
cd autoencoder-project
  • Confirm the folder was created by running:
ls

You should see an empty directory (no output is expected since it is brand new).

  • Press Cmd+Space (macOS) or the Windows key (Windows) to open your search bar.
  • Type VS Code and press Enter to open it.
  • In VS Code, click File in the top menu bar.
  • Click Open Folder.
  • Navigate to your Desktop, select the autoencoder-project folder, and click Open.
  • Create a new file by clicking File then New File.
  • Save it immediately with Cmd+S (macOS) or Ctrl+S (Windows) and name it autoencoder.py.

VS Code asking you to trust the folder?

If VS Code shows a trust dialog, click Yes, I trust the authors. This is normal for newly opened folders.

help me open my project in VS Code

Add imports and configuration constants

Every PyTorch project starts by importing the framework modules and defining the hyperparameters that control training. These constants live at the top of the file so they are easy to find and tweak later.

  • In autoencoder.py, add the following imports and configuration:
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import matplotlib.pyplot as plt

# --- Configuration ---
LATENT_DIM = 32
BATCH_SIZE = 64
EPOCHS = 10
LEARNING_RATE = 1e-3

What does this code do?

  • torch is the core PyTorch library for tensor operations and neural networks.
  • torch.nn provides the building blocks for defining network layers.
  • torch.optim contains optimizers like Adam that update the model's weights during training.
  • DataLoader handles batching and shuffling your training data.
  • torchvision gives you access to common image datasets and transforms.
  • matplotlib is for plotting training progress and visualizing results.

The four constants control your autoencoder's architecture and training:

  • LATENT_DIM = 32 means the encoder compresses each 784-pixel image down to just 32 numbers.
  • BATCH_SIZE = 64 means the model processes 64 images at a time during training.
  • EPOCHS = 10 is how many times the model sees the entire dataset.
  • LEARNING_RATE = 1e-3 controls how big each weight update is (smaller = more careful learning).
  • Save autoencoder.py with Cmd+S (macOS) or Ctrl+S (Windows).

Seeing red underlines in VS Code?

Red squiggly underlines on your import lines usually mean VS Code has not detected your Python environment. Click the Python version in the bottom-left status bar and select the interpreter where you installed PyTorch.

help me fix Python import errors in VS Code

✔️ Awesome, I've got everything!

Great. Double check you have saved your file.

ⓧ I'd like to double check the full code

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

# --- Configuration ---
LATENT_DIM = 32
BATCH_SIZE = 64
EPOCHS = 10
LEARNING_RATE = 1e-3

Your environment is ready and your project file is set up. Next up, you will load the MNIST dataset and see real handwritten digit images for the first time.

Load and Explore the MNIST Dataset

Your autoencoder needs data to learn from. Before the network can compress and reconstruct images, it needs thousands of examples to study.

In this step, you'll load the MNIST handwritten digit dataset and confirm it's working by visualizing a few sample images. MNIST contains 70,000 grayscale images of digits 0 through 9, each sized 28x28 pixels.

In this step, get ready to:
  • Download and load the MNIST training and test datasets.
  • Create DataLoaders to feed batches of images into your network.
  • Visualize sample images to confirm the data loaded correctly.
Download MNIST and create DataLoaders

MNIST is a classic machine learning dataset containing 60,000 training images and 10,000 test images of handwritten digits. Each image is a 28x28 pixel grayscale picture, which means each image is really just 784 numbers (28 times 28).

What is a DataLoader?

A DataLoader groups your images into batches (here, 64 images at a time) and feeds them to the network during training. Processing images in batches is more memory-efficient than loading all 60,000 at once.

  • In autoencoder.py, find the line that says LEARNING_RATE = 1e-3, then copy the following code below it:
# --- Load Data ---
transform = transforms.ToTensor()

train_dataset = datasets.MNIST(
    root='./data', train=True, download=True, transform=transform
)
test_dataset = datasets.MNIST(
    root='./data', train=False, download=True, transform=transform
)

train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False)

What does this code do?

transforms.ToTensor() converts each image into a PyTorch tensor and normalizes pixel values from 0-255 down to the 0-1 range. This makes training more stable.

datasets.MNIST() downloads the dataset to a ./data folder the first time you run it. The train=True flag grabs the 60,000-image training set, while train=False grabs the 10,000-image test set.

The training loader uses shuffle=True so the network sees images in a different order each epoch. The test loader uses shuffle=False so results are reproducible.

  • Save autoencoder.py.

Seeing a download error?

The first run downloads about 60 MB of data. Make sure you have a stable internet connection.

If you see a permissions error when creating the ./data folder, make sure you're running the script from inside your project folder.

Why is my MNIST dataset failing to download?

Visualize sample images from the dataset

Before training, it helps to see what the data actually looks like. You'll display a grid of 8 sample digits using matplotlib and print the dataset dimensions to your terminal.

  • Add the following code directly below the test_loader line you just wrote:
# --- Visualize Sample Data ---
sample_images, sample_labels = next(iter(train_loader))
fig, axes = plt.subplots(1, 8, figsize=(12, 2))
for i in range(8):
    axes[i].imshow(sample_images[i].squeeze(), cmap='gray')
    axes[i].set_title(f'Label: {sample_labels[i].item()}')
    axes[i].axis('off')
plt.tight_layout()
plt.savefig('sample_data.png')
plt.show()

print(f"Training set size: {len(train_dataset)}")
print(f"Test set size: {len(test_dataset)}")
print(f"Image shape: {sample_images[0].shape}")

What does this code do?

next(iter(train_loader)) grabs the first batch of 64 images. You then display 8 of them in a row with their labels above each one.

.squeeze() removes the single color channel dimension so matplotlib can render the image as 2D grayscale. The cmap='gray' argument tells matplotlib to display in grayscale rather than a color map.

plt.savefig('sample_data.png') saves the figure to your project folder so you can reference it later. The print statements confirm the dataset loaded with the expected sizes.

  • Save autoencoder.py.
  • Run the script by executing this command in your terminal:
python3 autoencoder.py
  • You should see a window pop up displaying 8 handwritten digit images in grayscale. Close the window to let the script finish.
  • In your terminal, you should see output like:
Training set size: 60000
Test set size: 10000
Image shape: torch.Size([1, 28, 28])

Don't see the image window?

Some environments don't support plt.show() pop-up windows. Check your project folder for sample_data.png instead. The image is saved there regardless of whether the window appears.

If you see a ModuleNotFoundError, double-check that you installed matplotlib in the same Python environment you're running the script from.

Why isn't my matplotlib plot showing up?

✔️ Awesome, I've got everything!

Great. Double check you've saved your file before moving on.

ⓧ I'd like to double check the full code

Here's the complete autoencoder.py file so far:

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

# --- Configuration ---
LATENT_DIM = 32
BATCH_SIZE = 64
EPOCHS = 10
LEARNING_RATE = 1e-3

# --- Load Data ---
transform = transforms.ToTensor()

train_dataset = datasets.MNIST(
    root='./data', train=True, download=True, transform=transform
)
test_dataset = datasets.MNIST(
    root='./data', train=False, download=True, transform=transform
)

train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False)

# --- Visualize Sample Data ---
sample_images, sample_labels = next(iter(train_loader))
fig, axes = plt.subplots(1, 8, figsize=(12, 2))
for i in range(8):
    axes[i].imshow(sample_images[i].squeeze(), cmap='gray')
    axes[i].set_title(f'Label: {sample_labels[i].item()}')
    axes[i].axis('off')
plt.tight_layout()
plt.savefig('sample_data.png')
plt.show()

print(f"Training set size: {len(train_dataset)}")
print(f"Test set size: {len(test_dataset)}")
print(f"Image shape: {sample_images[0].shape}")

Your data pipeline is ready. You've confirmed that 70,000 digit images are loaded, batched, and flowing correctly. Next up, you'll build the autoencoder architecture that learns to compress these 784-pixel images into just 32 numbers.

Build the Autoencoder Architecture

Your MNIST data is loaded and you can see what the handwritten digits look like.

Now it is time to build the neural network that will learn to compress those images. An autoencoder has two halves: an encoder that squeezes data down to a smaller representation, and a decoder that tries to rebuild the original from that compressed form. In this step, you will define both halves and wire them together in PyTorch.

In this step, get ready to:
  • Define an Autoencoder class with separate encoder and decoder networks.
  • Implement a forward pass that flattens, encodes, and decodes each image.
  • Instantiate the model and verify the architecture prints correctly.
Define the Autoencoder class

Each MNIST image is 28x28 pixels, which flattened is 784 numbers. The encoder's job is to squeeze those 784 numbers down to just 32. The decoder then takes those 32 numbers and expands them back to 784.

  • In autoencoder.py, scroll to the bottom of the file (below the print(f"Image shape: ... line) and add the following class definition:
# --- Define the Autoencoder ---
class Autoencoder(nn.Module):
    def __init__(self, latent_dim):
        super().__init__()
        self.encoder = nn.Sequential(
            nn.Linear(784, 128),
            nn.ReLU(),
            nn.Linear(128, 64),
            nn.ReLU(),
            nn.Linear(64, latent_dim),
            nn.ReLU()
        )
        self.decoder = nn.Sequential(
            nn.Linear(latent_dim, 64),
            nn.ReLU(),
            nn.Linear(64, 128),
            nn.ReLU(),
            nn.Linear(128, 784),
            nn.Sigmoid()
        )

What does this code do?

The class inherits from nn.Module, which is PyTorch's base class for all neural networks. This gives your model the ability to track its own parameters and switch between training and evaluation modes.

  • The encoder compresses 784 inputs down to 128, then 64, then 32 (your LATENT_DIM). Each layer uses ReLU activation to introduce non-linearity.
  • The decoder mirrors the encoder in reverse: 32 up to 64, then 128, then 784. The final Sigmoid activation keeps output values between 0 and 1, matching the pixel range of your input images.
  • nn.Sequential chains layers together so data passes through them in order automatically.

Why Sigmoid on the decoder output?

The transforms.ToTensor() transform you applied earlier normalizes pixel values to the [0, 1] range. The decoder's output must land in the same range so the loss function can meaningfully compare reconstructed pixels to originals. Sigmoid guarantees every output is between 0 and 1.

Add the forward method and create the model

The forward method defines how data flows through the network when you pass an image in. PyTorch calls this method automatically when you use the model.

  • Still inside the Autoencoder class (indented one level), add the forward method below the decoder. Then, outside the class, add the model instantiation:
    def forward(self, x):
        x = x.view(x.size(0), -1)
        latent = self.encoder(x)
        reconstructed = self.decoder(latent)
        return reconstructed


model = Autoencoder(latent_dim=LATENT_DIM)
print(model)

What does this code do?

  • x.view(x.size(0), -1) flattens each 28x28 image into a single row of 784 numbers. The -1 tells PyTorch to calculate the remaining dimension automatically.
  • The flattened input passes through the encoder (784 down to 32), then the decoder (32 back up to 784).
  • model = Autoencoder(latent_dim=LATENT_DIM) creates an instance of your network with a 32-dimensional latent space.
  • print(model) displays the full layer structure so you can verify everything is wired correctly.
  • Save autoencoder.py.
  • Run the script in your terminal:
python autoencoder.py
  • After the sample image window appears and you close it, you should see the model architecture printed in your terminal. It shows the encoder compressing from 784 down to 32, and the decoder expanding from 32 back to 784.

How many parameters does this model have?

The total is approximately 222,000 trainable parameters. Most come from the first encoder layer (784 x 128 = 100,352 weights plus 128 biases) and the last decoder layer (128 x 784 = 100,352 weights plus 784 biases). This is a small model that trains quickly on CPU.

Not seeing the model layers printed?

  • Check that model = Autoencoder(latent_dim=LATENT_DIM) and print(model) are outside the class (no indentation). They should be at the same level as the class keyword itself.
  • If you see an IndentationError, make sure the forward method is indented at the same level as __init__ (one tab or 4 spaces inside the class).
  • Make sure you close the sample image plot window before the script continues to print the model.

Help me debug my Autoencoder class definition in PyTorch

✔️ Awesome, I've got everything!

Great. Double check you have saved your file before moving on.

ⓧ I'd like to double check the full code

Here is the complete autoencoder.py so far:

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

# --- Configuration ---
LATENT_DIM = 32
BATCH_SIZE = 64
EPOCHS = 10
LEARNING_RATE = 1e-3

# --- Load Data ---
transform = transforms.ToTensor()

train_dataset = datasets.MNIST(
    root='./data', train=True, download=True, transform=transform
)
test_dataset = datasets.MNIST(
    root='./data', train=False, download=True, transform=transform
)

train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False)

# --- Visualize Sample Data ---
sample_images, sample_labels = next(iter(train_loader))
fig, axes = plt.subplots(1, 8, figsize=(12, 2))
for i in range(8):
    axes[i].imshow(sample_images[i].squeeze(), cmap='gray')
    axes[i].set_title(f'Label: {sample_labels[i].item()}')
    axes[i].axis('off')
plt.tight_layout()
plt.savefig('sample_data.png')
plt.show()

print(f"Training set size: {len(train_dataset)}")
print(f"Test set size: {len(test_dataset)}")
print(f"Image shape: {sample_images[0].shape}")


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

    def forward(self, x):
        x = x.view(x.size(0), -1)
        latent = self.encoder(x)
        reconstructed = self.decoder(latent)
        return reconstructed


model = Autoencoder(latent_dim=LATENT_DIM)
print(model)

Your autoencoder architecture is locked in. Next up, you will train this model on 60,000 digit images and watch the loss drop as it learns to compress and reconstruct them.

Train the Autoencoder

Your autoencoder architecture is defined. Now it's time to teach it how to compress and reconstruct images.

Training means showing the model thousands of images repeatedly, measuring how far off its reconstructions are, and adjusting its weights to get closer each time.

In this step, get ready to:
  • Set up the loss function and optimizer.
  • Run the training loop for 10 epochs.
  • Plot the training loss curve to confirm learning.
Set up the loss function and optimizer

Before the model can learn, it needs two things: a way to measure how wrong its reconstructions are, and a strategy for adjusting its weights to reduce that error.

  • In autoencoder.py, scroll down to the bottom of your file (below print(model)) and add the following code:
# --- Training ---
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)

train_losses = []

What does this code do?

  • nn.MSELoss() is the loss function. It measures the average squared difference between each original pixel and the reconstructed pixel. A lower loss means better reconstructions.
  • optim.Adam is the optimizer. It decides how to adjust the model's weights based on the loss. Adam adapts its learning speed for each parameter automatically.
  • train_losses is an empty list that will store the average loss for each epoch so you can plot it later.

You won't see any output yet. The next chunk uses these to train the model.

Write the training loop

The training loop is where the model actually learns. Each full pass through the training data is called an epoch.

Within each epoch, the model sees every batch of images, reconstructs them, measures its error, and updates its weights to improve.

  • Add the training loop directly below the train_losses = [] line:
for epoch in range(EPOCHS):
    model.train()
    epoch_loss = 0.0

    for batch_images, _ in train_loader:
        batch_images_flat = batch_images.view(batch_images.size(0), -1)

        optimizer.zero_grad()
        reconstructed = model(batch_images)
        loss = criterion(reconstructed, batch_images_flat)
        loss.backward()
        optimizer.step()

        epoch_loss += loss.item()

    avg_loss = epoch_loss / len(train_loader)
    train_losses.append(avg_loss)
    print(f"Epoch [{epoch + 1}/{EPOCHS}], Loss: {avg_loss:.6f}")

What does this code do?

Each epoch runs the same four-step cycle on every batch of 64 images:

  • batch_images_flat reshapes each 28x28 image into a flat vector of 784 pixels, which is what the loss function compares against.
  • optimizer.zero_grad() resets the gradient calculations from the previous batch.
  • reconstructed = model(batch_images) passes images through the encoder and decoder, producing reconstructions.
  • loss.backward() calculates how much each weight contributed to the error (backpropagation).
  • optimizer.step() nudges all weights in the direction that reduces the loss.

At the end of each epoch, the average loss is printed so you can watch the model improve in real time.

Why compare against the flat image?

The decoder outputs a 784-length vector (not a 28x28 grid). The loss function needs the original image in the same flat shape so it can compare pixel-by-pixel.

You won't see output until you run the script. The next substep adds the loss plot, then you'll run everything together.

Plot the training loss curve

A loss curve is the simplest way to confirm your model is learning. If the line goes down over epochs, the model is getting better at reconstructing images.

  • Add the plotting code below the training loop (after the print statement at the end of the loop):
# --- Plot Training Loss ---
plt.figure(figsize=(8, 4))
plt.plot(range(1, EPOCHS + 1), train_losses, marker='o')
plt.xlabel('Epoch')
plt.ylabel('MSE Loss')
plt.title('Training Loss Over Epochs')
plt.grid(True)
plt.savefig('training_loss.png')
plt.show()

What does this code do?

  • plt.plot() draws the loss values across all 10 epochs with circular markers at each point.
  • plt.savefig('training_loss.png') saves the plot as an image file in your project folder.
  • plt.show() opens the plot in a window so you can see it immediately.
  • Save autoencoder.py.
  • Run the script in your terminal by running:
python3 autoencoder.py
  • Close the sample data plot window when it appears (this is from the earlier code).
  • Watch the terminal. You should see 10 lines of output showing the loss decreasing from around 0.06 down to around 0.02.
  • After training finishes, the loss curve plot appears. You should see a clear downward trend from left to right.

What should I see?

Your terminal output will look something like this:

Epoch [1/10], Loss: 0.058432
Epoch [2/10], Loss: 0.034521
Epoch [3/10], Loss: 0.028876
...
Epoch [10/10], Loss: 0.019543

Your exact numbers will differ slightly because the weights start at random values. The key sign of success is that the loss consistently goes down.

Training takes 2-3 minutes

The model is processing 60,000 images ten times over on your CPU. This is normal. If you want faster training in future projects, look into GPU acceleration with CUDA.

Loss not decreasing?

  • Make sure optimizer.zero_grad() comes BEFORE the forward pass. Without it, gradients accumulate from previous batches and confuse the optimizer.
  • Check that your loss compares reconstructed against batch_images_flat (the flattened original), not batch_images (which is still 28x28).
  • Verify that loss.backward() comes before optimizer.step(). The optimizer needs the gradients calculated first.

My training loss stays flat or increases. Help me debug my training loop.

✔️ Awesome, I've got everything!

Great. Double check you've saved your file.

ⓧ I'd like to double check the full code

Here is the full autoencoder.py file up to this point:

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

# --- Configuration ---
LATENT_DIM = 32
BATCH_SIZE = 64
EPOCHS = 10
LEARNING_RATE = 1e-3

# --- Load Data ---
transform = transforms.ToTensor()

train_dataset = datasets.MNIST(
    root='./data', train=True, download=True, transform=transform
)
test_dataset = datasets.MNIST(
    root='./data', train=False, download=True, transform=transform
)

train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False)

# --- Visualize Sample Data ---
sample_images, sample_labels = next(iter(train_loader))
fig, axes = plt.subplots(1, 8, figsize=(12, 2))
for i in range(8):
    axes[i].imshow(sample_images[i].squeeze(), cmap='gray')
    axes[i].set_title(f'Label: {sample_labels[i].item()}')
    axes[i].axis('off')
plt.tight_layout()
plt.savefig('sample_data.png')
plt.show()

print(f"Training set size: {len(train_dataset)}")
print(f"Test set size: {len(test_dataset)}")
print(f"Image shape: {sample_images[0].shape}")


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

    def forward(self, x):
        x = x.view(x.size(0), -1)
        latent = self.encoder(x)
        reconstructed = self.decoder(latent)
        return reconstructed


model = Autoencoder(latent_dim=LATENT_DIM)
print(model)

# --- Training ---
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)

train_losses = []

for epoch in range(EPOCHS):
    model.train()
    epoch_loss = 0.0

    for batch_images, _ in train_loader:
        batch_images_flat = batch_images.view(batch_images.size(0), -1)

        optimizer.zero_grad()
        reconstructed = model(batch_images)
        loss = criterion(reconstructed, batch_images_flat)
        loss.backward()
        optimizer.step()

        epoch_loss += loss.item()

    avg_loss = epoch_loss / len(train_loader)
    train_losses.append(avg_loss)
    print(f"Epoch [{epoch + 1}/{EPOCHS}], Loss: {avg_loss:.6f}")

# --- Plot Training Loss ---
plt.figure(figsize=(8, 4))
plt.plot(range(1, EPOCHS + 1), train_losses, marker='o')
plt.xlabel('Epoch')
plt.ylabel('MSE Loss')
plt.title('Training Loss Over Epochs')
plt.grid(True)
plt.savefig('training_loss.png')
plt.show()

Your autoencoder has learned to compress and reconstruct images. Next up, you'll see exactly how well it does by comparing original images with their reconstructions and exploring the latent space.

Visualize Results and Explore the Latent Space

Your autoencoder has been trained for 10 epochs and the loss curve shows it learned to compress and reconstruct images. Now comes the exciting part: seeing what the network actually learned.

In this step, you will feed test images through the trained model and compare the reconstructions side by side with the originals. Then you will peek inside the encoder's latent space to see how it organizes different digits.

In this step, get ready to:
  • Reconstruct test images and display them alongside the originals.
  • Extract latent vectors for the entire test set.
  • Plot the latent space colored by digit label to visualize learned groupings.
Visualize reconstructed images

The whole point of an autoencoder is to reconstruct its input after compressing it. Comparing original and reconstructed images side by side is the most direct way to judge whether the model learned anything useful.

  • In autoencoder.py, scroll to the bottom of your file (below the training loss plot code).
  • Add the reconstruction visualization code by copying the following:
# --- Visualize Reconstructions ---
model.eval()

with torch.no_grad():
    test_images, _ = next(iter(test_loader))
    reconstructed = model(test_images)

n_images = 8
fig, axes = plt.subplots(2, n_images, figsize=(12, 3))

for i in range(n_images):
    axes[0, i].imshow(test_images[i].squeeze(), cmap='gray')
    axes[0, i].axis('off')
    if i == 0:
        axes[0, i].set_title('Original')

    axes[1, i].imshow(reconstructed[i].view(28, 28), cmap='gray')
    axes[1, i].axis('off')
    if i == 0:
        axes[1, i].set_title('Reconstructed')

plt.tight_layout()
plt.savefig('reconstructions.png')
plt.show()

What does this code do?

First, model.eval() switches the model from training mode to evaluation mode. This disables training-specific behaviors like dropout.

The torch.no_grad() context manager tells PyTorch not to track gradients, saving memory since we are only running inference.

The 2-row grid shows originals on top and reconstructions on the bottom. The reconstructed tensor comes out as a flat 784-length vector, so .view(28, 28) reshapes it back into image dimensions for display.

  • Save autoencoder.py and run it.
  • You should see a 2-row grid pop up with 8 original digit images on top and their reconstructions on the bottom. The reconstructions will look slightly blurry but clearly show the same digits.

Reconstructions look like random noise?

Make sure the training loop ran successfully before this code. Check your terminal output for decreasing loss values across epochs.

Verify that model.eval() appears before the torch.no_grad() block.

Why do my reconstructed images look like random static instead of digits?

Extract latent vectors from the test set

Your encoder compresses each 784-pixel image into just 32 numbers. By extracting these latent vectors for all 10,000 test images, you can visualize how the encoder organizes different digits in its compressed space.

  • Add the latent space extraction code below the reconstruction visualization:
# --- Explore Latent Space ---
model.eval()
latent_vectors = []
labels = []

with torch.no_grad():
    for images, targets in test_loader:
        encoded = model.encoder(images.view(images.size(0), -1))
        latent_vectors.append(encoded)
        labels.append(targets)

latent_vectors = torch.cat(latent_vectors, dim=0)
labels = torch.cat(labels, dim=0)

What does this code do?

Instead of running images through the full autoencoder, this code passes them through only model.encoder. This gives you the compressed 32-dimensional representation of each image.

The loop processes all test batches and collects their latent vectors and labels into lists. Then torch.cat concatenates them into single tensors of shape 10000x32 for the vectors and 10000 for the labels.

  • Save autoencoder.py.
  • You will not see a visual change yet. The next chunk uses these vectors to create the scatter plot.
Create a latent space scatter plot

The latent space has 32 dimensions, so you cannot visualize all of it at once. Plotting just the first 2 dimensions gives a partial window into how the encoder organizes data. If digits cluster together, the encoder has learned that similar digits share similar compressed features.

  • Add the scatter plot code below the latent vector extraction:
plt.figure(figsize=(8, 6))
scatter = plt.scatter(
    latent_vectors[:, 0].numpy(),
    latent_vectors[:, 1].numpy(),
    c=labels.numpy(),
    cmap='tab10',
    alpha=0.5,
    s=5
)
plt.colorbar(scatter, label='Digit')
plt.xlabel('Latent Dimension 1')
plt.ylabel('Latent Dimension 2')
plt.title('Latent Space Visualization (First 2 Dimensions)')
plt.savefig('latent_space.png')
plt.show()

What does this code do?

This creates a scatter plot where each point is one test image, positioned by its first two latent dimensions. The c=labels.numpy() argument colors each point by its digit class (0-9), and cmap='tab10' assigns a distinct color to each digit.

Points that cluster by color mean the encoder placed similar digits near each other in the compressed representation. This is a sign that meaningful structure was learned, not just random compression.

  • Save autoencoder.py and run the full script.
  • You should see a colorful scatter plot appear with some grouping of same-colored points. The plot saves as latent_space.png.

Why only 2 out of 32 dimensions?

Plotting 2 dimensions out of 32 is a simplification. You are seeing a slice of the full latent space, not the whole picture. Techniques like t-SNE or PCA can project all 32 dimensions into 2D more faithfully, but that is outside the scope of this project.

Even with this partial view, you should see some clustering. If the same-colored points tend to group together, it means those two latent dimensions carry digit-identity information.

Scatter plot looks completely random?

Check that you are using model.encoder (not the full model) when extracting latent vectors. Passing through the full model gives you reconstructed pixels, not latent vectors.

Confirm the images are being flattened with images.view(images.size(0), -1) before passing to the encoder.

Help me understand why my latent space scatter plot has no clustering.

✔️ Awesome, I've got everything!

Great. Double check you have saved autoencoder.py and that both reconstructions.png and latent_space.png appear in your project folder.

ⓧ I'd like to double check the full code

Here is the complete autoencoder.py file at this point:

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

# --- Configuration ---
LATENT_DIM = 32
BATCH_SIZE = 64
EPOCHS = 10
LEARNING_RATE = 1e-3

# --- Load Data ---
transform = transforms.ToTensor()

train_dataset = datasets.MNIST(
    root='./data', train=True, download=True, transform=transform
)
test_dataset = datasets.MNIST(
    root='./data', train=False, download=True, transform=transform
)

train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False)

# --- Visualize Sample Data ---
sample_images, sample_labels = next(iter(train_loader))
fig, axes = plt.subplots(1, 8, figsize=(12, 2))
for i in range(8):
    axes[i].imshow(sample_images[i].squeeze(), cmap='gray')
    axes[i].set_title(f'Label: {sample_labels[i].item()}')
    axes[i].axis('off')
plt.tight_layout()
plt.savefig('sample_data.png')
plt.show()

print(f"Training set size: {len(train_dataset)}")
print(f"Test set size: {len(test_dataset)}")
print(f"Image shape: {sample_images[0].shape}")


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

    def forward(self, x):
        x = x.view(x.size(0), -1)
        latent = self.encoder(x)
        reconstructed = self.decoder(latent)
        return reconstructed


model = Autoencoder(latent_dim=LATENT_DIM)
print(model)

# --- Training ---
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)

train_losses = []

for epoch in range(EPOCHS):
    model.train()
    epoch_loss = 0.0

    for batch_images, _ in train_loader:
        batch_images_flat = batch_images.view(batch_images.size(0), -1)

        optimizer.zero_grad()
        reconstructed = model(batch_images)
        loss = criterion(reconstructed, batch_images_flat)
        loss.backward()
        optimizer.step()

        epoch_loss += loss.item()

    avg_loss = epoch_loss / len(train_loader)
    train_losses.append(avg_loss)
    print(f"Epoch [{epoch + 1}/{EPOCHS}], Loss: {avg_loss:.6f}")

# --- Plot Training Loss ---
plt.figure(figsize=(8, 4))
plt.plot(range(1, EPOCHS + 1), train_losses, marker='o')
plt.xlabel('Epoch')
plt.ylabel('MSE Loss')
plt.title('Training Loss Over Epochs')
plt.grid(True)
plt.savefig('training_loss.png')
plt.show()

# --- Visualize Reconstructions ---
model.eval()

with torch.no_grad():
    test_images, _ = next(iter(test_loader))
    reconstructed = model(test_images)

n_images = 8
fig, axes = plt.subplots(2, n_images, figsize=(12, 3))

for i in range(n_images):
    axes[0, i].imshow(test_images[i].squeeze(), cmap='gray')
    axes[0, i].axis('off')
    if i == 0:
        axes[0, i].set_title('Original')

    axes[1, i].imshow(reconstructed[i].view(28, 28), cmap='gray')
    axes[1, i].axis('off')
    if i == 0:
        axes[1, i].set_title('Reconstructed')

plt.tight_layout()
plt.savefig('reconstructions.png')
plt.show()

# --- Explore Latent Space ---
model.eval()
latent_vectors = []
labels = []

with torch.no_grad():
    for images, targets in test_loader:
        encoded = model.encoder(images.view(images.size(0), -1))
        latent_vectors.append(encoded)
        labels.append(targets)

latent_vectors = torch.cat(latent_vectors, dim=0)
labels = torch.cat(labels, dim=0)

plt.figure(figsize=(8, 6))
scatter = plt.scatter(
    latent_vectors[:, 0].numpy(),
    latent_vectors[:, 1].numpy(),
    c=labels.numpy(),
    cmap='tab10',
    alpha=0.5,
    s=5
)
plt.colorbar(scatter, label='Digit')
plt.xlabel('Latent Dimension 1')
plt.ylabel('Latent Dimension 2')
plt.title('Latent Space Visualization (First 2 Dimensions)')
plt.savefig('latent_space.png')
plt.show()

Secret mission

Build a Denoising Autoencoder

Your autoencoder can compress and reconstruct images. But can it fix them? Transform your model into a denoising autoencoder that learns to remove noise from corrupted images, proving it has learned the true structure of digits rather than just copying pixels.

Clean Up Your Resources

Clean Up Your Resources

Decide whether to keep your resources running, pause them to come back later, or delete them entirely.

This project runs entirely locally, so there are no ongoing costs.

Resources you used:

  • A data/ folder containing the downloaded MNIST dataset (~60 MB).
  • Several .png image files showing training progress and results.
  • The autoencoder.py script.

Keep everything running

No action needed. Choose this if you're still actively building or experimenting with your autoencoder.

  • The data/ folder and all .png files are entirely local to your machine.
  • There are no cloud resources, subscriptions, or running processes to worry about.
  • The MNIST dataset in data/ takes up approximately 60 MB of disk space.

Pause - I'll come back to this later

Since everything is local files, there's nothing to pause. Your project will be exactly where you left it when you return.

  • If you want to reclaim ~60 MB of disk space without losing your code, delete only the data/ folder.
  • The MNIST dataset will re-download automatically the next time you run autoencoder.py.
  • Delete the data/ folder by running this from inside your project folder:
rm -rf data/

On Windows, run this instead:

rmdir /s /q data

Delete - I don't want to use this again

Remove all project files and start fresh. This deletes your code, the dataset, and all generated images.

  • Close autoencoder.py in VS Code.
  • In your terminal, navigate to the parent directory of your project folder (the folder that contains autoencoder.py).

On macOS or Linux, run:

cd ~/Desktop
rm -rf autoencoder-project/

On Windows, run:

cd ~\Desktop
rmdir /s /q autoencoder-project

Folder name may differ

Replace autoencoder-project/ with whatever you named your project folder in Step 1.

  • This removes autoencoder.py, the data/ folder, and all .png output files in one step.
  • To also remove the Python packages you installed, run:
pip uninstall torch torchvision matplotlib -y

Nice Work!

Nice Work!

Nice work! You just built a fully-connected autoencoder from scratch in PyTorch that compresses and reconstructs handwritten digit images.

You've learned how to:

  • Build a complete autoencoder from scratch with separate encoder and decoder components in PyTorch.
  • Train the model on 60,000 handwritten digit images and watch the loss decrease over 10 epochs.
  • Visualize that the network learned to compress 784-dimensional images into just 32 numbers and reconstruct them.
  • Explore the latent space and observe how the encoder organizes digits into meaningful clusters.
  • Secret Mission: Extended the autoencoder to denoise corrupted images, proving it learned structural features rather than just copying pixels.

Ready to quiz yourself?