MNIST Vanilla Autoencoder Example#

This example demonstrates how to train a fully connected autoencoder on the MNIST dataset using pyautoencoder.

Configuration and Utilities#

We begin by importing the required libraries, defining the main hyperparameters, and setting up deterministic random seeds to ensure reproducibility.

 1from __future__ import annotations
 2
 3from pathlib import Path
 4import time
 5import random
 6
 7import numpy as np
 8import torch
 9import torch.nn as nn
10from torch.utils.data import DataLoader
11from torchvision import datasets, transforms
12import matplotlib.pyplot as plt
13
14from pyautoencoder.vanilla import AE
15
16
17# ---------------- Config ---------------- #
18LATENT_DIM = 128
19NUM_EPOCHS = 30
20BATCH_SIZE = 128
21LEARNING_RATE = 1e-3
22NUM_RECON_COLS = 10
23
24DEVICE = torch.device("cpu")
25SEED = 1926
26
27# ---------------- Utils ----------------- #
28def set_seed(seed: int) -> None:
29    random.seed(seed)
30    np.random.seed(seed)
31    torch.manual_seed(seed)
32
33

Key elements in this section include:

  • LATENT_DIM, NUM_EPOCHS, BATCH_SIZE – the primary model and training hyperparameters,

  • DEVICE – the device used for training (the CPU),

  • set_seed() – helper function to provide reproducible runs.

Data Loading#

We prepare the MNIST training and test datasets using torchvision.datasets.MNIST. A simple ToTensor transform is used to map images to the \([0, 1]\) range.

1def make_dataloaders(batch_size: int) -> tuple[DataLoader, DataLoader]:
2    tfm = transforms.ToTensor()  # maps to [0,1]
3    train_dataset = datasets.MNIST("./data", train=True,  download=True, transform=tfm)
4    test_dataset  = datasets.MNIST("./data", train=False, download=True, transform=tfm)
5
6    train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
7    test_dataloader  = DataLoader(test_dataset,  batch_size=batch_size, shuffle=False)
8    return train_dataloader, test_dataloader

The make_dataloaders() function returns two torch.utils.data.DataLoader objects that supply batches to the training and evaluation loops.

Model Definition#

The autoencoder consists of a compact fully connected encoder and decoder. The encoder flattens the input and maps it into a latent representation of dimension LATENT_DIM; the decoder reconstructs the image from this latent vector.

 1def make_autoencoder(latent_dim: int) -> AE:
 2    encoder = nn.Sequential(
 3        nn.Flatten(),
 4        nn.Linear(28 * 28, 256),
 5        nn.ReLU(),
 6        nn.Linear(256, latent_dim),
 7    )
 8    decoder = nn.Sequential(
 9        nn.Linear(latent_dim, 256),
10        nn.ReLU(),
11        nn.Linear(256, 28 * 28),
12        nn.Unflatten(-1, (1, 28, 28)),  # keep last layer linear: the Bernoulli likelihood expects logits
13    )
14    model = AE(encoder=encoder, decoder=decoder)
15    model.build(input_sample=torch.randn(1, 1, 28, 28))
16    return model

Notes:

  • The last layer with parameters is linear and produces logits; nn.Unflatten then reshapes them to the input shape [B, 1, 28, 28].

  • The model is built through AE.build(), which runs a forward pass on a representative sample. This architecture has no lazy layers, so the call only checks that encoder and decoder fit together.

  • Reconstruction loss is computed via the AE.compute_loss() method, called with likelihood='bernoulli', so the decoder output is interpreted as logits.

Training Loop#

Training uses standard mini-batch gradient descent with the torch.optim.Adam optimizer. During each iteration, we compute the reconstruction log-likelihood and related diagnostics.

 1def train(
 2    model: AE,
 3    train_loader: DataLoader,
 4    num_epochs: int,
 5    lr: float = LEARNING_RATE,
 6) -> None:
 7    model.to(DEVICE)
 8    model.train()
 9    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
10
11    start_time = time.time()
12    for epoch in range(1, num_epochs + 1):
13        running_log_likelihood = torch.zeros((), device=DEVICE)
14        n = 0
15
16        for x, _ in train_loader:
17            x = x.to(DEVICE)
18            optimizer.zero_grad()
19
20            out = model(x)
21            loss_info = model.compute_loss(x, out, likelihood='bernoulli')
22            loss_info.objective.backward()
23            optimizer.step()
24
25            batch_size = x.size(0)
26            running_log_likelihood += loss_info.diagnostics['log_likelihood'] * batch_size
27            n += batch_size
28
29        avg_log_likelihood = running_log_likelihood.item() / n
30        elapsed = time.time() - start_time
31
32        print(
33            f"Epoch {epoch:3d}/{num_epochs}  "
34            f"LogLikelihood={avg_log_likelihood:.4f} | "
35            f"(elapsed {elapsed:.1f}s)"
36        )

For each batch:

  • The model produces latent codes and reconstructions: out = model(x).

  • AE.compute_loss() returns a LossResult containing:

    • objective – batch-mean negative log-likelihood (NLL) in nats,

    • diagnostics['log_likelihood'] – batch-mean log-likelihood (negative of objective).

  • The loss is backpropagated through the entire model.

  • The log-likelihood is accumulated over the epoch (weighted by batch size), and its per-sample average is printed at the end of each epoch.

Visualizing Reconstructions#

After training completes, we visualize a few randomly chosen test images alongside their reconstructions. The decoder outputs logits, so we apply a sigmoid to convert them to pixel intensities suitable for display.

 1@torch.no_grad()
 2def plot_test_reconstructions(
 3    model: AE,
 4    test_loader: DataLoader,
 5    latent_dim: int,
 6    num_cols: int = NUM_RECON_COLS,
 7    fname: str | None = None,
 8) -> None:
 9    model.eval()
10
11    x_batch, _ = next(iter(test_loader))
12    idx = torch.randperm(x_batch.size(0))[:num_cols]
13    x = x_batch[idx].to(DEVICE)
14
15    out = model(x)
16    x_hat = torch.sigmoid(out.x_hat)
17
18    x = x.cpu().numpy()
19    x_hat = x_hat.cpu().numpy()
20
21    fig, axes = plt.subplots(2, num_cols, figsize=(2.5 * num_cols, 4.5))
22    if num_cols == 1:
23        axes = np.array([axes])
24
25    axes[0, 0].set_ylabel("Original", fontsize=18)
26    axes[1, 0].set_ylabel("Reconstruction", fontsize=18)
27
28    for r in range(num_cols):
29        axes[0, r].imshow(x[r, 0], cmap="gray", vmin=0, vmax=1)
30        axes[0, r].set_xticks([])
31        axes[0, r].set_yticks([])
32        axes[1, r].imshow(x_hat[r, 0], cmap="gray", vmin=0, vmax=1)
33        axes[1, r].set_xticks([])
34        axes[1, r].set_yticks([])
35
36    plt.tight_layout()
37    out_path = fname or f"ae_mnist_test_recon_latent{latent_dim}.png"
38    Path(out_path).parent.mkdir(parents=True, exist_ok=True)
39    plt.savefig(out_path, dpi=200, bbox_inches="tight")
40    plt.close(fig)
41    print(f"Saved figure -> {out_path}")

Putting It All Together#

The main() function orchestrates the full example: seeding, dataset preparation, model creation, training, and saving a figure with test-set reconstructions.

 1def main() -> None:
 2    print(f"\n=== Training AE (latent_dim={LATENT_DIM}) for {NUM_EPOCHS} epochs ===")
 3    print(f"Using device: {DEVICE}")
 4
 5    set_seed(SEED)
 6    train_loader, test_loader = make_dataloaders(BATCH_SIZE)
 7
 8    model = make_autoencoder(LATENT_DIM)
 9    train(model, train_loader, NUM_EPOCHS)
10    plot_test_reconstructions(
11        model,
12        test_loader,
13        latent_dim=LATENT_DIM,
14        num_cols=NUM_RECON_COLS,
15        fname=f"ae_mnist_test_recon_latent{LATENT_DIM}.png",
16    )
17    print("Done.")

The script can also be executed directly:

1if __name__ == "__main__":
2    main()

Full Example Script#

For completeness, the entire example script is shown below:

mnist_ae.py#
  1from __future__ import annotations
  2
  3from pathlib import Path
  4import time
  5import random
  6
  7import numpy as np
  8import torch
  9import torch.nn as nn
 10from torch.utils.data import DataLoader
 11from torchvision import datasets, transforms
 12import matplotlib.pyplot as plt
 13
 14from pyautoencoder.vanilla import AE
 15
 16
 17# ---------------- Config ---------------- #
 18LATENT_DIM = 128
 19NUM_EPOCHS = 30
 20BATCH_SIZE = 128
 21LEARNING_RATE = 1e-3
 22NUM_RECON_COLS = 10
 23
 24DEVICE = torch.device("cpu")
 25SEED = 1926
 26
 27# ---------------- Utils ----------------- #
 28def set_seed(seed: int) -> None:
 29    random.seed(seed)
 30    np.random.seed(seed)
 31    torch.manual_seed(seed)
 32
 33
 34# ---------------- Data ------------------ #
 35def make_dataloaders(batch_size: int) -> tuple[DataLoader, DataLoader]:
 36    tfm = transforms.ToTensor()  # maps to [0,1]
 37    train_dataset = datasets.MNIST("./data", train=True,  download=True, transform=tfm)
 38    test_dataset  = datasets.MNIST("./data", train=False, download=True, transform=tfm)
 39
 40    train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
 41    test_dataloader  = DataLoader(test_dataset,  batch_size=batch_size, shuffle=False)
 42    return train_dataloader, test_dataloader
 43
 44
 45# ---------------- Model ----------------- #
 46def make_autoencoder(latent_dim: int) -> AE:
 47    encoder = nn.Sequential(
 48        nn.Flatten(),
 49        nn.Linear(28 * 28, 256),
 50        nn.ReLU(),
 51        nn.Linear(256, latent_dim),
 52    )
 53    decoder = nn.Sequential(
 54        nn.Linear(latent_dim, 256),
 55        nn.ReLU(),
 56        nn.Linear(256, 28 * 28),
 57        nn.Unflatten(-1, (1, 28, 28)),  # keep last layer linear: the Bernoulli likelihood expects logits
 58    )
 59    model = AE(encoder=encoder, decoder=decoder)
 60    model.build(input_sample=torch.randn(1, 1, 28, 28))
 61    return model
 62
 63# ---------------- Train ----------------- #
 64def train(
 65    model: AE,
 66    train_loader: DataLoader,
 67    num_epochs: int,
 68    lr: float = LEARNING_RATE,
 69) -> None:
 70    model.to(DEVICE)
 71    model.train()
 72    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
 73
 74    start_time = time.time()
 75    for epoch in range(1, num_epochs + 1):
 76        running_log_likelihood = torch.zeros((), device=DEVICE)
 77        n = 0
 78
 79        for x, _ in train_loader:
 80            x = x.to(DEVICE)
 81            optimizer.zero_grad()
 82
 83            out = model(x)
 84            loss_info = model.compute_loss(x, out, likelihood='bernoulli')
 85            loss_info.objective.backward()
 86            optimizer.step()
 87
 88            batch_size = x.size(0)
 89            running_log_likelihood += loss_info.diagnostics['log_likelihood'] * batch_size
 90            n += batch_size
 91
 92        avg_log_likelihood = running_log_likelihood.item() / n
 93        elapsed = time.time() - start_time
 94
 95        print(
 96            f"Epoch {epoch:3d}/{num_epochs}  "
 97            f"LogLikelihood={avg_log_likelihood:.4f} | "
 98            f"(elapsed {elapsed:.1f}s)"
 99        )
100
101
102# ---------------- Plot (TEST samples) ---------------- #
103@torch.no_grad()
104def plot_test_reconstructions(
105    model: AE,
106    test_loader: DataLoader,
107    latent_dim: int,
108    num_cols: int = NUM_RECON_COLS,
109    fname: str | None = None,
110) -> None:
111    model.eval()
112
113    x_batch, _ = next(iter(test_loader))
114    idx = torch.randperm(x_batch.size(0))[:num_cols]
115    x = x_batch[idx].to(DEVICE)
116
117    out = model(x)
118    x_hat = torch.sigmoid(out.x_hat)
119
120    x = x.cpu().numpy()
121    x_hat = x_hat.cpu().numpy()
122
123    fig, axes = plt.subplots(2, num_cols, figsize=(2.5 * num_cols, 4.5))
124    if num_cols == 1:
125        axes = np.array([axes])
126
127    axes[0, 0].set_ylabel("Original", fontsize=18)
128    axes[1, 0].set_ylabel("Reconstruction", fontsize=18)
129
130    for r in range(num_cols):
131        axes[0, r].imshow(x[r, 0], cmap="gray", vmin=0, vmax=1)
132        axes[0, r].set_xticks([])
133        axes[0, r].set_yticks([])
134        axes[1, r].imshow(x_hat[r, 0], cmap="gray", vmin=0, vmax=1)
135        axes[1, r].set_xticks([])
136        axes[1, r].set_yticks([])
137
138    plt.tight_layout()
139    out_path = fname or f"ae_mnist_test_recon_latent{latent_dim}.png"
140    Path(out_path).parent.mkdir(parents=True, exist_ok=True)
141    plt.savefig(out_path, dpi=200, bbox_inches="tight")
142    plt.close(fig)
143    print(f"Saved figure -> {out_path}")
144
145
146# ---------------- Main ------------------ #
147def main() -> None:
148    print(f"\n=== Training AE (latent_dim={LATENT_DIM}) for {NUM_EPOCHS} epochs ===")
149    print(f"Using device: {DEVICE}")
150
151    set_seed(SEED)
152    train_loader, test_loader = make_dataloaders(BATCH_SIZE)
153
154    model = make_autoencoder(LATENT_DIM)
155    train(model, train_loader, NUM_EPOCHS)
156    plot_test_reconstructions(
157        model,
158        test_loader,
159        latent_dim=LATENT_DIM,
160        num_cols=NUM_RECON_COLS,
161        fname=f"ae_mnist_test_recon_latent{LATENT_DIM}.png",
162    )
163    print("Done.")
164
165
166if __name__ == "__main__":
167    main()