.. _mnist_ae_example:

MNIST Vanilla Autoencoder Example
=================================

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

Configuration and Utilities
---------------------------

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

.. literalinclude:: ../../../examples/mnist_ae.py
   :language: python
   :linenos:
   :end-before: # ---------------- Data

Key elements in this section include:

* :data:`LATENT_DIM`, :data:`NUM_EPOCHS`, :data:`BATCH_SIZE` –
  the primary model and training hyperparameters,
* :data:`DEVICE` – the device used for training (the CPU),
* :func:`set_seed` – helper function to provide reproducible runs.


Data Loading
------------

We prepare the MNIST training and test datasets using
:class:`torchvision.datasets.MNIST`. A simple ``ToTensor`` transform is
used to map images to the :math:`[0, 1]` range.

.. literalinclude:: ../../../examples/mnist_ae.py
   :language: python
   :linenos:
   :pyobject: make_dataloaders

The :func:`make_dataloaders` function returns two
:class:`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 :data:`LATENT_DIM`; the decoder reconstructs
the image from this latent vector.

.. literalinclude:: ../../../examples/mnist_ae.py
   :language: python
   :linenos:
   :pyobject: make_autoencoder

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 :meth:`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 :meth:`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
:class:`torch.optim.Adam` optimizer. During each iteration, we compute
the reconstruction log-likelihood and related diagnostics.

.. literalinclude:: ../../../examples/mnist_ae.py
   :language: python
   :linenos:
   :pyobject: train

For each batch:

* The model produces latent codes and reconstructions: ``out = model(x)``.
* :meth:`AE.compute_loss` returns a :class:`~pyautoencoder.loss.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.

.. literalinclude:: ../../../examples/mnist_ae.py
   :language: python
   :linenos:
   :pyobject: plot_test_reconstructions


Putting It All Together
-----------------------

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

.. literalinclude:: ../../../examples/mnist_ae.py
   :language: python
   :linenos:
   :pyobject: main

The script can also be executed directly:

.. literalinclude:: ../../../examples/mnist_ae.py
   :language: python
   :linenos:
   :start-at: if __name__ == "__main__":


Full Example Script
-------------------

For completeness, the entire example script is shown below:

.. literalinclude:: ../../../examples/mnist_ae.py
   :language: python
   :linenos:
   :caption: mnist_ae.py