Package: normflows
normflows: A PyTorch Package for Normalizing Flowsnormflowsis a PyTorch implementation of discrete normalizing flows. Many popular flow architectures are implemented,
see thelist below. The package can be easilyinstalled via pip.
The basic usage is describedhere, and afull documentationis available as well. A more detailed description of this package is given in ouraccompanying paper.Several sample use cases are provided in theexamplesfolder,
includingGlow,
aVAE, and
aResidual Flow.
Moreover, two simple applications are highlighed in theexamples section. You can run them
yourself in Google Colab using the links below to get a feeling fornormflows.LinkDescriptionReal NVP applied to a 2D bimodal target distributionModeling a distribution on a cylinder surface with a neural spline flowModeling and generating CIFAR-10 images with GlowImplemented FlowsArchitectureReferencePlanar FlowRezende & Mohamed, 2015Radial FlowRezende & Mohamed, 2015NICEDinh et al., 2014Real NVPDinh et al., 2017GlowKingma et al., 2018Masked Autoregressive FlowPapamakarios et al., 2017Neural Spline FlowDurkan et al., 2019Circular Neural Spline FlowRezende et al., 2020Residual FlowChen et al., 2019Stochastic Normalizing FlowWu et al., 2020Note that Neural Spline Flows with circular and non-circular coordinates
are supported as well.InstallationThe latest version of the package can be installed via pippip install normflowsAt least Python 3.7 is required. If you want to use a GPU, make sure that
PyTorch is set up correctly by following the instructions at thePyTorch website.To run the example notebooks clone the repository firstgit clone https://github.com/VincentStimper/normalizing-flows.gitand then install the dependencies.pip install -r requirements_examples.txtUsageA normalizing flow consists of a base distribution, defined innf.distributions.base,
and a list of flows, given innf.flows.
Let's assume our target is a 2D distribution. We pick a diagonal Gaussian
base distribution, which is the most popular choice. Our flow shall be aReal NVP modeland, therefore, we need
to define a neural network for computing the parameters of the affine coupling
map. One dimension is used to compute the scale and shift parameter for the
other dimension. After each coupling layer we swap their roles.importnormflowsasnf# Define 2D Gaussian base distributionbase=nf.distributions.base.DiagGaussian(2)# Define list of flowsnum_layers=32flows=[]foriinrange(num_layers):# Neural network with two hidden layers having 64 units each# Last layer is initialized by zeros making training more stableparam_map=nf.nets.MLP([1,64,64,2],init_zeros=True)# Add flow layerflows.append(nf.flows.AffineCouplingBlock(param_map))# Swap dimensionsflows.append(nf.flows.Permute(2,mode='swap'))Once they are set up, we can define anf.NormalizingFlowmodel. If the target density is available, it can be added to the model
to be used during training. Sample target distributions are given innf.distributions.target.# If the target density is not givenmodel=nf.NormalizingFlow(base,flows)# If the target density is giventarget=nf.distributions.target.TwoMoons()model=nf.NormalizingFlow(base,flows,target)The loss can be computed with the methods of the model and minimized.# When doing maximum likelihood learning, i.e. minimizing the forward KLD# with no target distribution givenloss=model.forward_kld(x)# When minimizing the reverse KLD based on the given target distributionloss=model.reverse_kld(num_samples=512)# Optimization as usualloss.backward()optimizer.step()ExamplesWe provide several illustrative examples of how to use the package in theexamplesdirectory. Among them are implementations ofGlow,
aVAE, and
aResidual Flow.
More advanced experiments can be done with the scripts listed in therepository about resampled base distributions,
see itsexperimentsfolder.Below, we consider two simple 2D examples.Real NVP applied to a 2D bimodal target distributionInthis notebook,
which can directly be opened inColab,
we consider a 2D distribution with two half-moon-shaped modes as a target. We approximate it with a Real NVP model
and obtain the following results.Note that there might be a density filament connecting the two modes, which is due to an architectural limitation
of normalizing flows, especially prominent in Real NVP. You can find out more about it inthis paper.Modeling a distribution on a cylinder surface with a neural spline flowInanother example,
which is available inColabas well, we apply a Neural Spline Flow model to a distribution defined on a cylinder. The resulting density is visualized below.This example is considered in thepaperaccompanying this repository.SupportIf you have problems, please read thepackage documentationand check out theexamples sectionabove. You are also welcome tocreate issues on GitHubto get help. Note that it is
worthwhile browsing the existingopenandclosedissues, which might
address the problem you are facing.ContributingIf you find a bug or have a feature request, pleasefile an issue on GitHub.You are welcome to contribute to the package by fixing the bug or adding the feature yourself. If you want to
contribute, please add tests for the code you added or modified and ensure it passes successfully by runningpytest.
This can be done by simply executingpytestwithin your local version of the repository. Make sure you code is well documented, and we also encourage contributions
to the existing documentation. Once you finished coding and testing, pleasecreate a pull request on GitHub.Used byThe package has been used in several research papers. Some of them are listed below.Andrew Campbell, Wenlong Chen, Vincent Stimper, José Miguel Hernández-Lobato, and Yichuan Zhang.A gradient based strategy for Hamiltonian Monte Carlo hyperparameter optimization.
In Proceedings of the 38th International Conference on Machine Learning, pp. 1238–1248. PMLR, 2021.Code available on GitHub.Vincent Stimper, Bernhard Schölkopf, and José Miguel Hernández-Lobato.Resampling Base Distributions of Normalizing Flows.
In Proceedings of The 25th International Conference on Artificial Intelligence and Statistics, volume 151, pp. 4915–4936, 2022.Code available on GitHub.Laurence I. Midgley, Vincent Stimper, Gregor N. C. Simm, Bernhard Schölkopf, and José Miguel Hernández-Lobato.Flow Annealed Importance Sampling Bootstrap.
The Eleventh International Conference on Learning Representations, 2023.Code available on GitHub.Arnau Quera-Bofarull, Joel Dyer, Anisoara Calinescu, J. Doyne Farmer, and Michael Wooldridge.BlackBIRDS: Black-Box Inference foR Differentiable Simulators.
Journal of Open Source Software, 8(89), 5776, 2023.Code available on GitHub.Moreover, theboltzgenpackage
has been build uponnormflows.CitationIf you usenormflows, please cite thecorresponding paperas follows.Stimper et al., (2023). normflows: A PyTorch Package for Normalizing Flows.
Journal of Open Source Software, 8(86), 5361,https://doi.org/10.21105/joss.05361Bibtex@article{Stimper2023, 
  author = {Vincent Stimper and David Liu and Andrew Campbell and Vincent Berenz and Lukas Ryll and Bernhard Schölkopf and José Miguel Hernández-Lobato}, 
  title = {normflows: A PyTorch Package for Normalizing Flows}, 
  journal = {Journal of Open Source Software}, 
  volume = {8},
  number = {86}, 
  pages = {5361}, 
  publisher = {The Open Journal}, 
  doi = {10.21105/joss.05361}, 
  url = {https://doi.org/10.21105/joss.05361}, 
  year = {2023}
}