Quickstart

LatentLego is like a lego box with building blocks for autoencoders. These building blocks are hierarchically assembled into larger components, each of which can be used in a modular way to define a model. To make this more clear, lets start with some examples.

Using pre-defined models

LatentLego hosts a model zoo with a number of common autoencoder architectures ready to use. They are subclassed Keras Model objects and understand common methods like .fit() and .predict(). So the quickest way to train an autoencoder with LatentLego is simply

import numpy as np
from latent.models import Autoencoder

x_train = np.array([np.random.poisson(lam=5, size=100) for i in range(100)])

# Creates an autoencoder model with a 2-dimensional latent space
ae = Autoencoder(latent_dim=2, x_dim=x_train.shape[1], activation='relu')

# Compiled the model with a poisson loss and a Adam optimizer
ae.compile(loss='mse', optimizer='adam')

# Trains the model on x_train
ae.fit(x_train, epochs=10, batch_size=10)

This trained autoencoder model can now be used to yield a low dimensional representation of the input (or new) data. In line with the scikit-learn API, autoencoder models in LatentLego have a .transform() method, which returns the low dimensional representation.

x_latent = ae.transform(x_train)

Using .predict() gives us the reconstructred input data after being passed through the bottleneck. This representation can be used for denoising of single-cell data (Erslan 2019).

x_recon = ae.predict(x_train)

Defining models by combining encoder and decoder

The second way to define a model with LatentLego is to pick any of the provided encoder and decoder models and combine them using the Autoencoder data structure. For instance, if you want to create a variational autoencoder with a negative binomial reconstruction loss you can simply combine a VariationalEncoder with a NegativeBinomialDecoder:

from latent.modules import VariationalEncoder, NegativeBinomialDecoder

# Creates a VariationalEncoder with a standard normal prior
encoder = VariationalEncoder(latent_dim=20, prior='normal', kld_weight=0.01)
# Creates a NegativeBinomialDecoder with a constant dispersion estimate
decoder = NegativeBinomialDecoder(x_dim=x_train.shape[1], dispersion='constant')

# Constructs an Autoencoder object with predefined encoder and decoder
ae = Autoencoder(
    encoder=encoder,
    decoder=decoder
)
ae.compile()

# Fit input data using size factors for each cell
x_sf = np.array([1.]*100)
ae.fit([x_train, x_sf], epochs=10, batch_size=10)

You might have noticed that this time we did not define a loss function. This is because some encoder and decoder models take care of the model loss on their own. Generally in LatentLego, encoders are responsible for adding losses related to the latent space, such as the KLDivergence loss of a VAE and decoders add the reconstruction loss. If you anyway decide to pass a loss function to .compile() it will be added to the reconstruction loss defined by the decoder.

The Autoencoder data structure subclasses the Keras Model and takes care of providing additional methods like .transform() and handling cases with multiple inputs like conditions or size factors. But since encoder and decoder models are just Keras Model objects themselves, they can also be used with the Keras Functional API:

from tensorflow.keras import Input, Model
from latent.modules import Encoder, Decoder

x = Input(shape=(x_train.shape[1],))
x_latent = Encoder(latent_dim=2)(x)

# Decoders take the latent and the original data space as the input
# so they can add the reconstruction loss
x_rec = Decoder(x_dim=x_train.shape[1])([x, x_latent])
ae = Model(inputs=x, outputs=x_rec)

ae.compile(loss='mse', optimizer='adam', run_eagerly=False)
ae.fit(x=x_train, y=x_train, epochs=10, batch_size=10)

Building autoencoders from scratch

Lastly, if you want to implement a model that can not be directly assembled in LatentLego, you may want to implement the encoder/decoder or both from scratch. For this, we also provide some convenient lower-level components that extend the reportoire of Keras and TensorFlow. For instance, the DenseStack layer strings together a sequence of DenseBlock layers and take care of batch normalization, dropout and conditional injection. We can combine it with some TensorFlow Probability magic and stack it on top of a NegativeBinomialDecoder to obtain a VAE with negative binomial loss.

import tensorflow as tf
from tensorflow.keras import Sequential, Input, Model
import tensorflow_probability as tfp
tfpl = tfp.layers
tfd = tfp.distributions
from latent.layers import DenseStack
from latent.modules import NegativeBinomialDecoder

latent_dim = 2
# Initiate standard normal prior
prior = tfd.Independent(tfd.Normal(loc=tf.zeros(latent_dim), scale=1),
    reinterpreted_batch_ndims=1)

# Construct encoder with sequential model
encoder = Sequential([
    DenseStack(hidden_units=[256,128,56], batchnorm=True, dropout_rate=0.2),
    tf.keras.layers.Dense(
        tfpl.MultivariateNormalTriL.params_size(latent_dim),
        activation='linear'),
    tfpl.MultivariateNormalTriL(
        latent_dim,
        activity_regularizer=tfpl.KLDivergenceRegularizer(prior))
])
# Initiate a NB decoder with a per-gene dispersion estimate
decoder = NegativeBinomialDecoder(x_dim=x_train.shape[1], dispersion='gene')

x_in = Input(shape=(x_train.shape[1],))
sf_in = Input(shape=(1,))
x_latent = encoder(x_in)
x_rec = decoder([x_in, x_latent, sf_in])
nbvae = Model([x_in, sf_in], x_rec)

# We specify no loss because NB loss is handled by the decoder
nbvae.compile(loss=None, optimizer='sgd')
nbvae.fit(x=[x_train, x_sf], y=x_train, epochs=10, batch_size=10)