Package: disent
🧶 DisentA modular disentangled representation learning framework built with PyTorch LightningVisit thedocsfor more info, or browse thereleases.Contributionsare welcome!────────────────NOTE:My MSc. research has movedhereSome of the contributions have been incorporated directly into disent────────────────Table Of ContentsOverviewFeaturesDatasetsFrameworksMetricsSchedules & AnnealingArchitectureExamplesPython ExampleHydra Config ExampleInstallDevelopmentWhy?OverviewDisent is a modular disentangled representation learning framework for auto-encoders,
built upon PyTorch-Lightning. This framework consists of various composable components
that can be used to build and benchmark various disentanglement vision tasks.The name of the framework is derived from bothdisentanglement and scientificdissent.Get started with disent by installing it with $pip install disentor cloning this repository.GoalsDisent aims to fill the following criteria:Providehigh quality,readable,consistentandeasily comparableimplementations of frameworksHighlight differencebetween framework implementations by overridinghooksand minimising duplicate codeUsebest practiceeg.torch.distributionsBe extremelyflexible& configurableSupport low memory systemsCiting DisentPlease use the following citation if you use Disent in your own research:@Misc{Michlo2021Disent,author={Nathan Juraj Michlo},title={Disent - A modular disentangled representation learning framework for pytorch},howpublished={Github},year={2021},url={https://github.com/nmichlo/disent}}FeaturesDisent includes implementations of modules, metrics and
datasets from various papers.Note that "🧵" means that the dataset, framework or metric was introduced by disent!DatasetsVarious common datasets used in disentanglement research are included with disent. The dataset loaders provide various features including:automatic downloads & preperationprepare=Trueautomatic hash verificationautomatic optimization of underlying hdf5 formats for
low-memory disk-based access.Data input and target dataset augmentations and transforms are supported, as well as augmentations
on the GPU or CPU at different points in the pipeline.Ground Truth:🚗Cars3D◻️dSprites🔺MPI3D🐘SmallNORB🌈Shapes3D🏹Sprites (custom)🧵 dSpritesImagenet:Version of DSprite with foreground or background deterministically masked out with tiny-imagenet data.Ground Truth Synthetic:🧵 XYSquares:Three non-overlapping squares that can move around a grid. This dataset is adversarial to VAEs that use pixel-wise reconstruction losses.🧵 XYObject:A simplistic version of dSprites with a single square.🧵 XYObjectShaded:Exact same dataset as XYObject, but ground truth factors have a different representation.FrameworksDisent provides the following Auto-Encoders and Variational Auto-Encoders!Unsupervised:AE:Auto-EncoderVAE: Variational Auto-EncoderBeta-VAE: VAE with Scaled LossDFC-VAE: Deep Feature Consistent VAEDIP-VAE: Disentangled Inferred Prior VAEInfoVAE: Information Maximizing VAEBetaTCVAE: Total Correlation VAEWeakly Supervised:Ada-GVAE: Adaptive GVAE,AdaVae.cfg(average_mode='gvae'), usually better than below!Ada-ML-VAE: Adaptive ML-VAE,AdaVae.cfg(average_mode='ml-vae')Supervised:TAE:Triplet Auto-EncoderTVAE: Triplet Variational Auto-EncoderIntroduced in DisentUnsupervised:🧵 Ada-TVAE-D: Adaptive Triplet VAE that uses data distances instead of ground-truth distances as the supervision signal.🧵 Ada-TAE-D:  Adaptive Triplet AE that uses data distances instead of ground-truth distances as the supervision signal.Weakly Supervised:🧵 Ada-AE: Adaptive AE, the auto-encoder version of the Ada-GVAESupervised:🧵 Ada-TVAE: Adaptive Triplet VAE, disentangled version of the TVAE🧵 Ada-TAE:  Adaptive Triplet AE, disentangled version of the TAE🏗 Todo:Many popular disentanglement frameworks still need to be added, please
submit an issue if you have a request for an additional framework.FactorVAEGroupVAEMLVAEMetricsVarious metrics are provided by disent that can be used to evaluate the
learnt representations of models that have been trained on ground-truth data.Disentanglement:FactorVAE ScoreDCIMIGSAPUnsupervised Scores🧵 Flatness Components:Measures of the three components needed to learn factored representations from distances. VAEs often learn the first two (correlation & linearity), and the can happen accidentally (axis-alignment)!🪡 Ground-Truth Correlation -The spearman rank correlation between latent distances and ground-truth distances.🪡 Linearity Ratio -How well factor traversals lie along an n-dimensional arbitrarily rotated line in the latent space🪡 Axis-Alignment Ratio -How well factor traversals are represented by a single latent variable, ie. an n-dimensional line that is axis-aligned.🧵 Flatness Score -Measuring the max distance between factor traversal embeddings and the path length of their embeddings.🏗 Todo:Some popular metrics still need to be added, please submit an issue if you wish to
add your own, or you have a request.DCIMIGModularity and ExplicitnessSchedules & AnnealingHyper-parameter annealing is supported through the use of schedules.
The currently implemented schedules include:Linear ScheduleCyclicScheduleCosine Wave ScheduleVarious other wrapper schedulesArchitectureThe disent module structure:disent.dataset: dataset wrappers, datasets & sampling strategiesdisent.dataset.data: raw datasetsdisent.dataset.sampling: sampling strategies forDisentDatasetwhen multiple elements are required by frameworks, eg. for triplet lossdisent.dataset.transform: common data transforms and augmentationsdisent.dataset.wrapper: wrapped datasets are no longer ground-truth datasets, these may have some elements masked out. We can still unwrap these classes to obtain the original datasets for benchmarking.disent.frameworks: frameworks, including Auto-Encoders and VAEsdisent.frameworks.ae: Auto-Encoder based frameworksdisent.frameworks.vae: Variational Auto-Encoder based frameworksdisent.metrics: metrics for evaluating disentanglement using ground truth datasetsdisent.model: common encoder and decoder models used for VAE researchdisent.nn: torch components for building models including layers, transforms, losses and general mathsdisent.schedule: annealing schedules that can be registered to a frameworkdisent.util: helper classes, functions, callbacks, anything unrelated to a pytorch system/model/framework.⚠️ The API IsMostlyStable ⚠️Disent is still under development. Features and APIs are subject to change!
However, I will try and minimise the impact of these.A small suite of tests currently exist which will be expanded upon in time.Hydra Experiment DirectoriesEasily run experiments with hydra config, these files
are not available frompip install.experiment/run.py: entrypoint for running basic experiments withhydraconfigexperiment/config/config.yaml: main configuration file, this is probably what you want to edit!experiment/config: root folder forhydraconfig filesexperiment/util: various helper code for experimentsExtending The Default ConfigsAll configs inexperiment/configcan easily be extended or overridden
without modifying any files. We can add a new config folder to the hydra search path
by setting the environment variableDISENT_CONFIGS_PREPENDto point to a config folder
that should take priority over those contained in the default folder.The advantage of this is that new frameworks and datasets can be used with experiments without cloning or modifying
disent itself. You can separate your research code from the library!See the examples in the docs for more information!ExamplesPython ExampleThe following is a basic working example of disent that trains a BetaVAE with a cyclic
beta schedule and evaluates the trained model with various metrics.💾 Basic ExampleimportlightningasLimporttorchfromtorch.utils.dataimportDataLoaderfromdisent.datasetimportDisentDatasetfromdisent.dataset.dataimportXYObjectDatafromdisent.dataset.samplingimportSingleSamplerfromdisent.dataset.transformimportToImgTensorF32fromdisent.frameworks.vaeimportBetaVaefromdisent.metricsimportmetric_dcifromdisent.metricsimportmetric_migfromdisent.modelimportAutoEncoderfromdisent.model.aeimportDecoderConv64fromdisent.model.aeimportEncoderConv64fromdisent.scheduleimportCyclicSchedule# create the dataset & dataloaders# - ToImgTensorF32 transforms images from numpy arrays to tensors and performs checks# - if you use `num_workers != 0` in the DataLoader, the make sure to#   wrap `trainer.fit` with `if __name__ == '__main__': ...`data=XYObjectData()dataset=DisentDataset(dataset=data,sampler=SingleSampler(),transform=ToImgTensorF32())dataloader=DataLoader(dataset=dataset,batch_size=128,shuffle=True,num_workers=0)# create the BetaVAE model# - adjusting the beta, learning rate, and representation size.module=BetaVae(model=AutoEncoder(# z_multiplier is needed to output mu & logvar when parameterising normal distributionencoder=EncoderConv64(x_shape=data.x_shape,z_size=10,z_multiplier=2),decoder=DecoderConv64(x_shape=data.x_shape,z_size=10),),cfg=BetaVae.cfg(optimizer='adam',optimizer_kwargs=dict(lr=1e-3),loss_reduction='mean_sum',beta=4,))# cyclic schedule for target 'beta' in the config/cfg. The initial value from the# config is saved and multiplied by the ratio from the schedule on each step.# - based on: https://arxiv.org/abs/1903.10145module.register_schedule('beta',CyclicSchedule(period=1024,# repeat every: trainer.global_step % period))# train model# - for 2048 batches/stepstrainer=L.Trainer(max_steps=2048,gpus=1iftorch.cuda.is_available()elseNone,logger=False,enable_checkpointing=False)trainer.fit(module,dataloader)# compute disentanglement metrics# - we cannot guarantee which device the representation is on# - this will take a while to runget_repr=lambdax:module.encode(x.to(module.device))metrics={**metric_dci(dataset,get_repr,num_train=1000,num_test=500,show_progress=True),**metric_mig(dataset,get_repr,num_train=2000),}# evaluateprint('metrics:',metrics)Visit thedocsfor more examples!Hydra Config ExampleThe entrypoint for basic experiments isexperiment/run.py.Some configuration will be required, but basic experiments can
be adjusted by modifying theHydra Config 1.1files inexperiment/config.Modifying the mainexperiment/config/config.yamlis all you
need for most basic experiments. The main config file contains
a defaults list with entries corresponding to yaml configuration
files (config options) in the subfolders (config groups) inexperiment/config/<config_group>/<option>.yaml.💾 Config Defaults Exampledefaults:# data-sampling:default__bb-dataset:xyobject-augment:none# system-framework:adavae_os-model:vae_conv64# training-optimizer:adam-schedule:beta_cyclic-metrics:fast-run_length:short# logs-run_callbacks:vis-run_logging:wandb# runtime-run_location:local-run_launcher:local-run_action:train# <rest of config.yaml left out>...Easily modify  any of these values to adjust how the basic experiment
will be run. For example, changeframework: adavaetoframework: betavae, or
change the dataset fromxyobjecttoshapes3d. Add new options by adding new
yaml files in the config group folders.Weights and Biasesis supported by changingrun_logging: nonetorun_logging: wandb. However, you will need to login from the command line. W&B logging supports
visualisations of latent traversals.InstallpipinstalldisentOtherwise, to install from source we recommend using a conda virtual environment.⤵️ Install from Source# clone the repogitclonehttps://github.com/nmichlo/disentcddisent# create and activate the conda environment [py38,py39,py310]condacreate-ndisent-py310python=3.10
condaactivatedisent-py310# check that the correct python version is usedwhichpython
whichpip# make sure to upgrade pippipinstall--upgradepip# install minimal requirementspipinstall-rrequirements.txt# (optional) install extra requirements# - first do the above because torch is required to compile torchsort while installingpipinstall-rrequirements-extra.txt# (optional) install test requirementspipinstall-rrequirements-test.txtDevelopmentMake sure to installpre-commithooks to ensure code is automatically formatted
correctly when committing or pushing changes todisent.# install git hookspipinstallpre-commit
pre-commitinstall# manually trigger all pre-commit hookspre-commitrun--all-filesTo run tests locally, make sure to install all the test and extra dependencies in your
environment.pipinstall-rrequirements.txt# torchsort first requires torch to be installedpipinstall-rrequirements-extra.txt-rrequirements-test.txtWhy?Created as part of my Computer Science MSc which ended early 2022.I needed custom high quality implementations of various VAE's.A pytorch version ofdisentanglement_lib.I didn't have time to wait forWeakly-Supervised Disentanglement Without Compromisesto release
their code as part of disentanglement_lib. (As of September 2020 it has been released, but has unresolveddiscrepencies).disentanglement_lib still uses outdated Tensorflow 1.0, and the flow of data is unintuitive because of its use ofGin Config.