🎓 A1Learner
🔍

PyTorch Autoencoder Basics

An autoencoder is a model that squeezes each input into a smaller code and rebuilds it from that code. This lesson squeezes 2D points into a 1-number code and checks what the rebuild keeps.

The squeeze and the rebuild

The encoder is the squeezing half of the model, and the decoder is the rebuilding half. One point goes in as two numbers, sits in the code as one number, and comes back as two numbers.

torch.manual_seed(1)
encoder = nn.Sequential(nn.Linear(2, 16), nn.ReLU(), nn.Linear(16, 1))
decoder = nn.Sequential(nn.Linear(1, 16), nn.ReLU(), nn.Linear(16, 2))
point = torch.tensor([[2.0, 2.0]])
code = encoder(point)
print(code.shape)
print(decoder(code).shape)

The rebuild error is one number

Mean squared error scores the rebuild against the original point by point. A small score means the rebuild sits close to the original.

pred = torch.tensor([[2.4, 2.3], [-2.3, -2.5]])
true = torch.tensor([[2.5, 2.2], [-2.4, -2.6]])
print(round(nn.functional.mse_loss(pred, true).item(), 4))

Training compares rebuilds with originals

Each pass scores the rebuilds against the same points that went in. The model gets no other target.

for epoch in range(300):
    loss = loss_fn(model(points), points)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

Example

import torch
from torch import nn

torch.manual_seed(0)
per_blob = 50
blob_a = torch.randn(per_blob, 2) * 0.5 + torch.tensor([-2.0, -2.0])
blob_b = torch.randn(per_blob, 2) * 0.5 + torch.tensor([2.0, 2.0])
points = torch.cat([blob_a, blob_b])


class Autoencoder(nn.Module):
    def __init__(self):
        super().__init__()
        self.encoder = nn.Sequential(nn.Linear(2, 16), nn.ReLU(), nn.Linear(16, 1))
        self.decoder = nn.Sequential(nn.Linear(1, 16), nn.ReLU(), nn.Linear(16, 2))

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


torch.manual_seed(1)
model = Autoencoder()
optimizer = torch.optim.Adam(model.parameters(), lr=0.02)
loss_fn = nn.MSELoss()

for epoch in range(300):
    loss = loss_fn(model(points), points)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

with torch.no_grad():
    rebuilt = model(points)

print("Final loss:", round(loss.item(), 4))
print("Original -> rebuilt:")
for i in [0, 10, 25, 50, 75, 99]:
    o = points[i].tolist()
    r = rebuilt[i].tolist()
    print(f"({o[0]:.2f}, {o[1]:.2f}) -> ({r[0]:.2f}, {r[1]:.2f})")

Key points

  • An autoencoder learns its own code by rebuilding its input
  • Use nn.MSELoss to score a rebuild against its original
  • A 1-number code can keep only one fact about each point
  • The rebuilt points land near the line between the two blob centres

In this example, the seed builds two blobs of 50 points around (-2, -2) and (2, 2). A second seed builds the autoencoder, so its starting weights repeat on every run. The encoder squeezes each point through a hidden layer of 16 into one number. The decoder rebuilds two numbers from that one number. Adam trains the model for 300 passes against the points themselves. The final loss prints as 0.0876. Six originals then print next to their rebuilds. Each rebuild sits close to its original and close to the line between the centres. The sideways detail of each point is what the 1-number code drops.

Note

Nothing labels the code, so its meaning appears only in the rebuilds

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

What does a 1-number bottleneck force an autoencoder to do?