flowcon 0.0.4

Normalizing flows in PyTorch. An extension of nflows.

FlowConductor: (Conditional) Normalizing Flows and
bijective Layers for Pytorch

Documentation

About

FlowConductor provides a collection of normalizing flows in PyTorch. It's core logic and transformations were originally based on the nflows package. The main focus lies in implementing more flow layers from the literature in one consistent framework, and adding support for conditional normalizing flows. In the original nflows package, conditional networks were restricted to using a conditional base distribution. In FlowConductor, nearly every layer can be conditional :). In particular, we support conditional transformations based on hypernetworks.

While using the package, we extensively expanded it, implementing layers such as invertible residual networks, and focusing on conditional normalizing flows. In addition, we improved the testing, so that the Jacobian determinants are actually reliably compared to a reference based on brute-force autograd.

The bijective layers we additionally provide includes but are not limited to Planar Flows, (conditional) invertible ResNets/DenseNets, a variant of neural autoregressive flows, and a basic support of continuous normalizing flows (and FFJORD) based on the torchdiffeq package.

Install

PIP

FlowConductor is installable via pip. We recommend using a virtual environment, where you set up your pytorch version beforehand. You can check out in ./docker which pytorch versions we test for, but in general there shouldn't be any complications for any version after 1.13.

You may either install the latest release from pipy:

$  pip install flowcon

or install it directly from github via pip

$  pip install git+https://github.com/FabricioArendTorres/FlowConductor.git

Of course, you may also just download the repo and install it locally

$ git clone https://github.com/FabricioArendTorres/FlowConductor
$ cd FlowConductor
$ pip install . 

Docker

We provide some basic Dockerfiles in ./docker, which are very simple extensions of the pytorch docker images. The dockerfiles we list are the ones used for testing, so you can be sure they work. If you are unfamiliar with Docker, you can use our package with it as follows (assuming it is at least installed).

This also works on Windows (cpu at least)!

$ git clone https://github.com/FabricioArendTorres/FlowConductor
$ cd FlowConductor

# Build the docker image, see the ./docker dir for different versions.
$ docker build -f ./docker/Dockerfile-pytorchlatest -t flowc-pytorchlatest .

# you can run the tests with
docker run flowc-pytorchlatest pytest /flowc

For working with this container, you may either choose to adapt our Dockerfiles, or simply bind the current directory when starting the container interactively. For the latter, you can run a script (here examples/toy_2d.py) with

$ docker run --rm -it -v .:/app flowc-pytorchlatest python examples/toy_2d.py Or you may swap an interactive shell within the container with

$ docker run --rm -it -v .:/app flowc-pytorchlatest
$ python examples/toy_2d.py

Package Usage

As the core is based on nflows, its usage is similar. To define a flow:

from flowcon import transforms, distributions, flows

# Define an invertible transformation.
transform = transforms.CompositeTransform([
  transforms.MaskedAffineAutoregressiveTransform(features=2, hidden_features=4),
  transforms.RandomPermutation(features=2)
])

# Define a base distribution.
base_distribution = distributions.StandardNormal(shape=[2])

# Combine into a flow.
flow = flows.Flow(transform=transform, distribution=base_distribution)

To evaluate log probabilities of inputs:

log_prob = flow.log_prob(inputs)

To sample from the flow:

samples = flow.sample(num_samples)

Additional examples of the workflow are provided in examples folder.

Changes and added features compared to nflows

The core logic of the code for LFlows (i.e. the nflows/ directory) is based on the nflows package. Aside from added features, our extension provides tests for the calculation of the Jacobian log-determinant, which is at the heart of Normalizing Flow.

Added Layers / Flwos:

• (Conditional) Sum-of-Sigmoid Layers
• Cholesky Outer Product for flows on symmetric positive definite matrices
• Lipschitz Constrained invertible DenseNets In particular, we provide three ways to condition these of these transformations without affecting the invertibility.
• Transformations for which the inverse is only known to exist, but not available:
  • (Conditional) Planar Flow
  • (Conditional) Sylvester Flow
• Conditional Versions of existing non-conditional transformations from nflows. Can be found for imports at nflows.transforms.conditional.*:
  • LU Transform
  • Orthogonal Transforms based on parameterized Householder projections
  • SVD based on the Orthogonal transforms
  • Shift Transform
• Conditional Versions of existing auto-regressive Variations, i.e. getting rid of the autoregressive parts.
  • ConditionalPiecewiseRationalQuadraticTransform
  • ConditionalUMNNTransform