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/pyautoencoders.git
cd pyautoencoders
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),
)

# 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.randn(32, 1, 28, 28)
output = model(x_batch)
loss_results = model.compute_loss(x_batch, output, likelihood='bernoulli')
optimizer.zero_grad()
loss_results.objective.backward()
optimizer.step()

# encode/decode run under inference_mode automatically
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),
)

# Create and build the model
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.randn(32, 1, 28, 28)
output = model(x_batch, S=1)  # S=1: single Monte Carlo sample
loss_results = model.compute_loss(x_batch, output, beta=1, likelihood='bernoulli')
optimizer.zero_grad()
loss_results.objective.backward()
optimizer.step()

Core Concepts

ModelOutput

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

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

Loss Result

Loss functions return structured results with optimization targets and diagnostics:

loss_results = loss_fn(x, model_output)
loss_results.objective.backward()      # Optimize this
print(loss_results.diagnostics)        # Log these for additional diagnostics