Getting Started

pyautoencoder is a PyTorch library for building and training state-of-the-art autoencoders. It is designed to host a growing collection of autoencoder variants, with structured outputs and modular loss functions to make experimentation easy.

Installation

Install pyautoencoder via pip:

pip install pyautoencoder

Or install from source:

git clone https://github.com/andrea-pollastro/pyautoencoder.git
cd pyautoencoder
pip install -e .

Quick Start

Vanilla Autoencoder (AE)

Train a simple autoencoder on MNIST:

import torch
import torch.nn as nn
from pyautoencoder.vanilla import AE

# Define encoder and decoder
encoder = nn.Sequential(
    nn.Flatten(),
    nn.Linear(784, 256),
    nn.ReLU(),
    nn.Linear(256, 64),
)

decoder = nn.Sequential(
    nn.Linear(64, 256),
    nn.ReLU(),
    nn.Linear(256, 784),
    nn.Unflatten(-1, (1, 28, 28)),  # x_hat must have the same shape as x
)

# Create and build the model
model = AE(encoder=encoder, decoder=decoder)
model.build(torch.randn(1, 1, 28, 28))

# Training loop
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

x_batch = torch.rand(32, 1, 28, 28)  # stand-in for an MNIST batch, values in [0, 1]
output = model(x_batch)
loss_results = model.compute_loss(x_batch, output, likelihood='bernoulli')
optimizer.zero_grad()
loss_results.objective.backward()
optimizer.step()

# Inference: encode/decode track no gradients but keep the train/eval mode
model.eval()
z = model.encode(x_batch).z       # AEEncodeOutput → latent tensor
x_hat = model.decode(z).x_hat     # AEDecodeOutput → reconstruction tensor

Variational Autoencoder (VAE)

Train a VAE with the reparameterization trick:

import torch
import torch.nn as nn
from pyautoencoder.variational import VAE

# Define encoder and decoder
encoder = nn.Sequential(
    nn.Flatten(),
    nn.Linear(784, 256),
    nn.ReLU(),
    nn.Linear(256, 128),
)

decoder = nn.Sequential(
    nn.Linear(64, 256),
    nn.ReLU(),
    nn.Linear(256, 784),
    nn.Unflatten(-1, (1, 28, 28)),  # x_hat must have the same shape as x
)

# Create and build the model (the encoder must output flat features [B, F])
model = VAE(encoder=encoder, decoder=decoder, latent_dim=64)
model.build(torch.randn(1, 1, 28, 28))

# Training loop
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

x_batch = torch.rand(32, 1, 28, 28)  # stand-in for an MNIST batch, values in [0, 1]
output = model(x_batch, S=1)  # S=1: single Monte Carlo sample, x_hat is [32, 1, 1, 28, 28]
loss_results = model.compute_loss(x_batch, output, beta=1, likelihood='bernoulli')
optimizer.zero_grad()
loss_results.objective.backward()
optimizer.step()

Core Concepts

Model outputs

All forward passes return structured output dataclasses that are easy to inspect:

output = model(x)  # the AE above
print(output)
# AEOutput(x_hat=Tensor(shape=(32, 1, 28, 28), dtype=torch.float32),
#          z=Tensor(shape=(32, 64), dtype=torch.float32))

Loss Result

Each model’s compute_loss returns a LossResult with the objective to optimize and a dictionary of diagnostics:

loss_results = model.compute_loss(x, output)
loss_results.objective.backward()      # Optimize this
# Diagnostics are detached 0-dim tensors: call .item() when logging
print({name: val.item() for name, val in loss_results.diagnostics.items()})