Implementation:LaurentMazare Tch rs VAE Example

Knowledge Sources	LaurentMazare_Tch_rs
Domains	Generative Models, Deep Learning, Computer Vision
Last Updated	2026-02-08 00:00 GMT

Overview

Implements a Variational Autoencoder (VAE) on the MNIST dataset for learning a latent representation and generating handwritten digit images.

Description

This module implements a Variational Autoencoder following the architecture from the PyTorch VAE example. The Vae struct contains five fully connected layers arranged as an encoder-decoder pair.

The encoder maps 784-dimensional flattened MNIST images through a hidden layer of 400 units (with ReLU) to produce two 20-dimensional vectors: mu (mean) and logvar (log-variance) of the latent distribution. The reparameterization trick samples from this distribution by computing mu + eps * exp(0.5 * logvar) where eps is drawn from a standard normal distribution.

The decoder maps 20-dimensional latent vectors through a 400-unit hidden layer (ReLU) to produce 784-dimensional output with sigmoid activation, representing pixel probabilities.

The loss function combines two terms: binary cross-entropy (BCE) for reconstruction quality (comparing reconstructed pixels to originals), and KL divergence measuring how close the learned latent distribution is to a standard normal prior, following Appendix B of the original VAE paper (Kingma and Welling, 2014). Both terms are summed over all elements and the batch dimension.

Training runs for 20 epochs using Adam with learning rate 1e-3 on batches of 128 images. After each epoch, the model generates a grid of 64 sample images from random latent vectors, arranged in an 8x8 matrix and saved as s_{epoch}.png.

Usage

Use this implementation to experiment with variational autoencoders on MNIST. It requires the MNIST dataset files extracted in the data/ directory. The model trains on CUDA if available and saves generated sample grids to the current directory after each epoch.

Code Reference

Source Location

• Repository: LaurentMazare_Tch_rs
• File: examples/vae/main.rs
• Lines: 1-95

Signature

struct Vae {
    fc1: nn::Linear,   // 784 -> 400 (encoder hidden)
    fc21: nn::Linear,  // 400 -> 20  (mu)
    fc22: nn::Linear,  // 400 -> 20  (logvar)
    fc3: nn::Linear,   // 20 -> 400  (decoder hidden)
    fc4: nn::Linear,   // 400 -> 784 (decoder output)
}

impl Vae {
    fn new(vs: &nn::Path) -> Self
    fn encode(&self, xs: &Tensor) -> (Tensor, Tensor)
    fn decode(&self, zs: &Tensor) -> Tensor
    fn forward(&self, xs: &Tensor) -> (Tensor, Tensor, Tensor)
}

fn loss(recon_x: &Tensor, x: &Tensor, mu: &Tensor, logvar: &Tensor) -> Tensor
fn image_matrix(imgs: &Tensor, sz: i64) -> Result<Tensor>

pub fn main() -> Result<()>

Import

use anyhow::Result;
use tch::{nn, nn::Module, nn::OptimizerConfig, Kind, Reduction, Tensor};

I/O Contract

Input	Type	Description
xs (encode)	&Tensor	Batch of MNIST images, shape [batch, 1, 28, 28] or [batch, 784]
zs (decode)	&Tensor	Latent vectors, shape [batch, 20]
MNIST data	Files	train-images, train-labels, t10k-images, t10k-labels in data/

Output	Type	Description
encode	(Tensor, Tensor)	(mu, logvar), each shape [batch, 20]
decode	Tensor	Reconstructed images, shape [batch, 784] with sigmoid activation
forward	(Tensor, Tensor, Tensor)	(reconstruction, mu, logvar)
loss	Tensor	Scalar combining BCE reconstruction loss and KL divergence
image_matrix	Tensor	Grid image tensor, shape [1, rows*28, cols*28]
Sample images	PNG files	s_1.png through s_20.png (8x8 grids of generated digits)

Usage Examples

// Run VAE training on MNIST (entry point)
// Requires MNIST data files in data/ directory
vae::main()?;

// Internal usage during training:
let device = tch::Device::cuda_if_available();
let m = tch::vision::mnist::load_dir("data")?;
let vs = nn::VarStore::new(device);
let vae = Vae::new(&vs.root());
let mut opt = nn::Adam::default().build(&vs, 1e-3)?;

// Training step
let (recon_batch, mu, logvar) = vae.forward(&bimages);
let loss = loss(&recon_batch, &bimages, &mu, &logvar);
opt.backward_step(&loss);

// Generate new samples from random latent vectors
let z = Tensor::randn([64, 20], tch::kind::FLOAT_CPU).to(device);
let generated = vae.decode(&z).view([64, 1, 28, 28]);

Related Pages

• Principle:LaurentMazare_Tch_rs_Variational_Autoencoder