Package: tfae
TensorFlow Autoencoders (TFAE)This package is aTensorFlowextension for building and training different types of autoencoders.Table of contentsMotivationCurrent packageInstallationUsageAPI referenceContributionMotivationAutoencoders – 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. Evenofficial tutorialon 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.Current packageTFAE 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 β-annealingand others.InstallationTFAE can be installed directly fromPyPI:pip install tfaeTensorFlow is the only dependency. Python 3.8-3.10 is required.UsageLet's take a quick example.Here we build and train shallow variational autoencoder regularized with KL-divergence:importtensorflowastffromtfae.modelsimportAutoencoderfromtfae.bottlenecksimportGaussianBottleneckfromtfae.regularizersimportGaussianKLDRegularizermodel=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 autoencoderGenerating handwritten digits with Convolutional Variational Autoencoderor explore API in the following section.API referenceTFAE includes:bottlenecksmodelsregularizersschedulersBottlenecksBottlenecks are layers placed in the middle of a model, connecting encoder with decoder.Every bottleneck extendsBaseBottleneck, which in turn extendstf.keras.layers.Layer.Vanilla bottleneckVanillaBottleneckit's a "semantic stub" for a dense layer – the most simple autoencoder bottleneck.Variational bottlenecksVariationalBottleneckis 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 parametersModelsTFAE includes two subclasses oftf.keras.Model:AutoencoderandDeepAutoencoder.classDiagram
    `tf.keras.Model` <|-- Autoencoder
    Autoencoder <|-- DeepAutoencoder

    Autoencoder : BaseBottleneck bottleneck
    Autoencoder : make_encoder()
    Autoencoder : make_decoder()
    DeepAutoencoder : Callable add_hiddenAutoencoderAutoencoderrepresents a simplest form of autoencoder with only one hidden layer as a bottleneck.See usage example.DeepAutoencoderDeepAutoencoderextendsAutoencoderand allows to build deeper models in a functional way: it'sadd_hiddenmethod constructs additional hidden layers.Let's take a quick example howadd_hiddenworks.It takes four parameters:input layernumber of current layershape of the input layerdimensionality of the bottleneckAnd 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:fromtfae.modelsimportDeepAutoencoderdefadd_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=2divisor=(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)returnx,layer_num==number_of_hidden_layersmodel=DeepAutoencoder(bottleneck=...add_hidden=add_hidden,)Custom modelsCustom models can be made by extendingAutoencoderclass.See an example.RegularizersIt often proves useful to regularize bottleneck, so encoder could learn better and disentangled representation.TFAE includes:L2RegularizerforVanillaBottleneckGaussianKLDRegularizerandGaussianReversedKLDRegularizerforGaussianBottleneckEvery TFAE regularizer extendsBaseRegularizer, which contains propertybeta: float– regularization factor:classDiagram
    `tf.keras.regularizers.Regularizer` <|-- BaseRegularizer
    BaseRegularizer <|-- L2Regularizer
    BaseRegularizer <|-- GaussianKLDRegularizer
    BaseRegularizer <|-- GaussianReversedKLDRegularizer

    BaseRegularizer: float betaA custom regulirizer can be applied by extendingBaseRegularizer.SchedulersRecent papers has shown that constant regularization factor can be an obstacle on the way to the better latent representation:β-VAEFixing Broken ELBOCyclical Annealing Scheduleβ-annealed VAEAll these papers are suggesting to vary regularization factor – let's call itβ– over time.TFAE containsDASRSchedulerwhich can handle different schedules covering the aforementioned papers.Every scheduler extendsBaseSchedulerwhich extendstf.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 cyclesDASRScheduler"DASR" stands for Delay, Attack, Sustain, Release.Let's say in "phase 1" we want to keepβ=0for first 5 epochs, then gradually rise it up toβ=1for 10 more epochs. In "phase 2" we'll keepβ=1until loss stops improving:fromtfae.modelsimportAutoencoderfromtfae.bottlenecksimportGaussianBottleneckfromtfae.regularizersimportGaussianKLDRegularizerfromtfae.schedulersimportDASRScheduler# 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 forcyclical β-annealing:scheduler=DASRScheduler(start_value=0.0,end_value=1.0,attack=10,sustain=10,cycles=4,)ContributionFeel free to createissuesor openpull-requests.