Normalizing Flows

Read the accompianing blog post

Normalizing flows have become popular for modeling distributions of data, for example unsupervised learning of image datasets.

For my first foray into Normalizing Flows I followed this great tutorial, which was originally written in Tensorflow 1. This repo is my implementation of the modern normalizing flow examples from the tutorial in tensorflow 2.

Installation

# 1. Clone the repo and create a virtual environment
git clone https://github.com/bsaund/normalizing_flows.git
cd normalizing_flows
python3 -m venv .venv
source .venv/bin/activate

# 2. Install dependencies
pip install --upgrade pip
pip install -r requirements.txt

# 3. Train and run the model (~1 hour on CPU; reduce train_iters in settings to speed up)
python normalizing_flows.py

GPU note: The default install uses tensorflow[and-cuda] (requires NVIDIA driver 520+ and CUDA 12). For CPU-only, replace tensorflow[and-cuda] with tensorflow-cpu in requirements.txt. Do not install both — they conflict.

Keras compatibility note: TF 2.16+ ships with Keras 3, which breaks TensorFlow Probability. tensorflow-probability[tf] automatically installs tf-keras (Keras 2) to work around this.

You should first see the training points:

After training you should see the learned mapping from a normal distribution to the training samples

What is going on?

normalizing_flows.py implements RealNVP. The image shows the transformation of samples from a 2D Gaussian through each of the layers of the network. The quadrants of the Gaussian are color-coded to visualize how the gaussian transforms.

It is composed 8 repeated units of RealNVP layers and Permutations, with occasional batch normaliziation bijectors.

Honestly, tensorflow_distributions does all of the heavy-lifting by implementing RealNVP layers, which take care of splitting the data and (most importantly) computing the jacobian needed to compute the gradient.

        for i in range(settings['num_bijectors']):
            self.bijector_fns.append(tfp.bijectors.real_nvp_default_template(hidden_layers=[512,512]))
            bijectors.append(
                tfb.RealNVP(num_masked=self.num_masked,
                            shift_and_log_scale_fn=self.bijector_fns[-1])
            )
            if i%2 == 0:
                bijectors.append(tfb.BatchNormalization())
            bijectors.append(tfb.Permute(permutation=[1,0]))

What to do next?

• Try making your own pictures. I just wrote "BRAD" in google sheets and exported it as .png.
• Try changing the number of layers, types of layers, or other settings
• Implement a normalizing flow on some image training data

Citation:

Feel free to copy and use this code in your projects. If you would like to cite this repo:

@misc{bsaund_2020_flow,
    author = {Brad Saund},
    title = {Normalizing Flows},
    year = 2020,
    month = {april},
    publisher = {Github},
    journal = {GitHub repository},
    url = {https://github.com/bsaund/normalizing_flows}
}