This package is a TensorFlow extension for building and training different types of autoencoders.
Autoencoders – AE – are widely used in various tasks: from computer vision to recommender systems and many others.
Tensorflow, being a flexible tool, does not provide AE-specific tools though.
The lack of such tools results into redundant code. Even official tutorial on this topic includes explicit looping through training steps and manual computation of losses and gradients.
Code like this is difficult to maintain.
TFAE makes process of building and training AEs easy yet flexible, keeping traditional TensorFlow style.
TFAE is a set of extensions for standard Tensorflow classes, such as Layers and Models. These can be used to build a wide range of models: shallow and deep, classical and variational.
In addition to layers and models, it contains extensions for Regularizers and Callbacks, which are used to control the training process, such as β-VAE, Cyclical β-annealing and others.
TFAE can be installed directly from PyPI:
pip install tfaeTensorFlow is the only dependency. Python 3.8-3.10 is required.
Let's take a quick example.
Here we build and train shallow variational autoencoder regularized with KL-divergence:
import tensorflow as tf
from tfae.models import Autoencoder
from tfae.bottlenecks import GaussianBottleneck
from tfae.regularizers import GaussianKLDRegularizer
model = Autoencoder(
bottleneck=GaussianBottleneck(
latent_dim=32,
kernel_regularizer=GaussianKLDRegularizer(),
),
)
model.compile(...)
model.fit(...)Now we can use this model to encode data:
encoded = model.encoder.predict(...)Or generate new samples:
generated = model.decoder.predict(...)Of course, it's possible to build deeper models and sophisticate training process.
Check out more examples:
- Search similar movies using autoencoder
- Generating handwritten digits with Convolutional Variational Autoencoder
or explore API in the following section.
TFAE includes:
Bottlenecks are layers placed in the middle of a model, connecting encoder with decoder.
Every bottleneck extends BaseBottleneck, which in turn extends tf.keras.layers.Layer.
VanillaBottleneck it's a "semantic stub" for a dense layer – the most simple autoencoder bottleneck.
VariationalBottleneck is a subclass for building model implementing variational inference. TFAE includes two bottlenecks for variational inference: Gaussian and Bernoulli:
classDiagram
`tf.keras.layers.Layer` <|-- BaseBottleneck
BaseBottleneck <|-- VanillaBottleneck
BaseBottleneck <|-- VariationalBottleneck
VariationalBottleneck <|-- GaussianBottleneck
VariationalBottleneck <|-- BernoulliBottleneck
BaseBottleneck : int latent_dim
VariationalBottleneck : int parameters
TFAE includes two subclasses of tf.keras.Model: Autoencoder and DeepAutoencoder.
classDiagram
`tf.keras.Model` <|-- Autoencoder
Autoencoder <|-- DeepAutoencoder
Autoencoder : BaseBottleneck bottleneck
Autoencoder : make_encoder()
Autoencoder : make_decoder()
DeepAutoencoder : Callable add_hidden
Autoencoder represents a simplest form of autoencoder with only one hidden layer as a bottleneck. See usage example.
DeepAutoencoder extends Autoencoder and allows to build deeper models in a functional way: it's add_hidden method constructs additional hidden layers.
Let's take a quick example how add_hidden works.
It takes four parameters:
- input layer
- number of current layer
- shape of the input layer
- dimensionality of the bottleneck
And returns a tuple of a new layer and a boolean indicating that current layer is the last.
This method is applied to both, encoder and decoder (but for decoder in a "mirror manner").
The following example demostrates how to create encoder and decoder with two hidden layers each. And both have a pyramidal structure:
from tfae.models import DeepAutoencoder
def add_hidden(
x: tf.keras.layers.Layer,
layer_num: int,
input_shape: tf.TensorShape,
latent_dim: int,
) -> tuple[tf.keras.layers.Layer, bool]:
number_of_hidden_layers = 2
divisor = (latent_dim / input_shape[-1]) ** (layer_num / (number_of_hidden_layers + 1))
units = int(divisor * input_shape[-1])
x = tf.keras.layers.Dense(units)(x)
return x, layer_num == number_of_hidden_layers
model = DeepAutoencoder(
bottleneck=...
add_hidden=add_hidden,
)Custom models can be made by extending Autoencoder class. See an example.
It often proves useful to regularize bottleneck, so encoder could learn better and disentangled representation.
TFAE includes:
L2RegularizerforVanillaBottleneckGaussianKLDRegularizerandGaussianReversedKLDRegularizerforGaussianBottleneck
Every TFAE regularizer extends BaseRegularizer, which contains property beta: float – regularization factor:
classDiagram
`tf.keras.regularizers.Regularizer` <|-- BaseRegularizer
BaseRegularizer <|-- L2Regularizer
BaseRegularizer <|-- GaussianKLDRegularizer
BaseRegularizer <|-- GaussianReversedKLDRegularizer
BaseRegularizer: float beta
A custom regulirizer can be applied by extending BaseRegularizer.
Recent papers has shown that constant regularization factor can be an obstacle on the way to the better latent representation:
All these papers are suggesting to vary regularization factor – let's call it β – over time.
TFAE contains DASRScheduler which can handle different schedules covering the aforementioned papers.
Every scheduler extends BaseScheduler which extends tf.keras.callbacks.Callback:
classDiagram
`tf.keras.callbacks.Callback` <|-- BaseScheduler
BaseScheduler <|-- DASRScheduler
BaseScheduler: calc()
DASRScheduler: float start_value
DASRScheduler: float end_value
DASRScheduler: int delay
DASRScheduler: int attack
DASRScheduler: int sustain
DASRScheduler: int release
DASRScheduler: int cycles
"DASR" stands for Delay, Attack, Sustain, Release.
Let's say in "phase 1" we want to keep β=0 for first 5 epochs, then gradually rise it up to β=1 for 10 more epochs. In "phase 2" we'll keep β=1 until loss stops improving:
from tfae.models import Autoencoder
from tfae.bottlenecks import GaussianBottleneck
from tfae.regularizers import GaussianKLDRegularizer
from tfae.schedulers import DASRScheduler
# Creating scheduler which will keep β=0 for 5 epochs
# and then gradually raise it up to β=1 for 10 more epochs:
scheduler = DASRScheduler(
start_value=0.0,
end_value=1.0,
delay=5,
attack=10,
)
# Note how we pass scheduler.value to the regularizer:
model = Autoencoder(
bottleneck=GaussianBottleneck(
latent_dim=32,
kernel_regulirizer=GaussianKLDRegularizer(
beta=scheduler.value,
)
)
)
model.compile(...)
# Phase 1.
#
# Scheduler have auto-calculated attribute "duration"
# which tells how many epochs it takes
# to go through all scheduled values of β.
#
# We also pass scheduler as a callback
# so so that he can be updated:
model.fit(
...
epochs=scheduler.duration,
callbacks=[
scheduler,
],
)
# Phase 2.
#
# Here we continue training until loss stops improving:
model.fit(
...,
initial_epoch=scheduler.duration,
callbacks=[
tf.keras.callbacks.EarlyStopping(...),
]
)Let's take one more example and implement schedule for cyclical β-annealing:
scheduler = DASRScheduler(
start_value=0.0,
end_value=1.0,
attack=10,
sustain=10,
cycles=4,
)Feel free to create issues or open pull requests.