Skip to content

slimgroup/InvertibleNetworks.jl

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

InvertibleNetworks.jl

Documentation Build Status
CI DOI

Building blocks for invertible neural networks in the Julia programming language.

  • Memory efficient building blocks for invertible neural networks
  • Hand-derived gradients, Jacobians $J$ , and $\log |J|$
  • Flux integration
  • Support for Zygote and ChainRules
  • GPU support
  • Includes various examples of invertible neural networks, normalizing flows, variational inference, and uncertainty quantification

Installation

InvertibleNetworks is registered and can be added like any standard Julia package with the command:

] add InvertibleNetworks

Uncertainty-aware image reconstruction

Due to its memory scaling InvertibleNetworks.jl has been particularily successful at Bayesian posterior sampling with simulation-based inference. To get started with this application please refer to a simple example (Conditional sampling for MNSIT inpainting) but please use this on your application and please reach out to us if you run into any trouble.

mnist_sampling_cond

Building blocks

  • 1x1 Convolutions using Householder transformations (example)

  • Residual block (example)

  • Invertible coupling layer from Dinh et al. (2017) (example)

  • Invertible hyperbolic layer from Lensink et al. (2019) (example)

  • Invertible coupling layer from Putzky and Welling (2019) (example)

  • Invertible recursive coupling layer HINT from Kruse et al. (2020) (example)

  • Activation normalization (Kingma and Dhariwal, 2018) (example)

  • Various activation functions (Sigmoid, ReLU, leaky ReLU, GaLU)

  • Objective and misfit functions (mean squared error, log-likelihood)

  • Dimensionality manipulation: squeeze/unsqueeze (column, patch, checkerboard), split/cat

  • Squeeze/unsqueeze using the wavelet transform

Examples

  • Invertible recurrent inference machines (Putzky and Welling, 2019) (generic example)

  • Generative models with maximum likelihood via the change of variable formula (example)

  • Glow: Generative flow with invertible 1x1 convolutions (Kingma and Dhariwal, 2018) (generic example, source)

GPU support

GPU support is supported via Flux/CuArray. To use the GPU, move the input and the network layer to GPU via |> gpu

using InvertibleNetworks, Flux

# Input
nx = 64
ny = 64
k = 10
batchsize = 4

# Input image: nx x ny x k x batchsize
X = randn(Float32, nx, ny, k, batchsize) |> gpu

# Activation normalization
AN = ActNorm(k; logdet=true) |> gpu

# Test invertibility
Y_, logdet = AN.forward(X)

Reference

If you use InvertibleNetworks.jl in your research, we would be grateful if you cite us with the following bibtex:

@article{orozco2023invertiblenetworks,
  title={InvertibleNetworks. jl: A Julia package for scalable normalizing flows},
  author={Orozco, Rafael and Witte, Philipp and Louboutin, Mathias and Siahkoohi, Ali and Rizzuti, Gabrio and Peters, Bas and Herrmann, Felix J},
  journal={arXiv preprint arXiv:2312.13480},
  year={2023}
}

Papers

The following publications use InvertibleNetworks.jl:

Authors

  • Rafael Orozco, Georgia Institute of Technology [rorozco@gatech.edu]

  • Philipp Witte, Georgia Institute of Technology (now Microsoft)

  • Gabrio Rizzuti, Utrecht University

  • Mathias Louboutin, Georgia Institute of Technology

  • Ali Siahkoohi, Georgia Institute of Technology

Acknowledgments

This package uses functions from NNlib.jl, Flux.jl and Wavelets.jl