*release*

.gitignore

@@ -0,0 +1,4 @@
+*.pyc
+*__pycache__*
+data/*
+pretrained_models

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2019 Ricky Tian Qi Chen
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -0,0 +1,87 @@
+# Residual Flows for Invertible Generative Modeling [[arxiv](https://arxiv.org/abs/1906.02735)]
+
+<p align="center">
+<img align="middle" src="./assets/flow_comparison.jpg" width="666" />
+</p>
+
+Building on the use of [Invertible Residual Networks](https://arxiv.org/abs/1811.00995) in generative modeling, we propose:
++ Unbiased estimation of the log-density of samples.
++ Memory-efficient reformulation of the gradients.
++ LipSwish activation function.
+
+As a result, Residual Flows scale to much larger networks and datasets.
+
+<p align="center">
+<img align="middle" src="./assets/celebahq_resflow.jpg" width="512" />
+</p>
+
+## Requirements
+
+ - PyTorch 1.0+
+ - Python 3.6+
+
+## Preprocessing
+ImageNet:
+1. Follow instructions in `preprocessing/create_imagenet_benchmark_datasets`.
+2. Convert .npy files to .pth using `preprocessing/convert_to_pth`.
+3. Place in `data/imagenet32` and `data/imagenet64`.
+
+CelebAHQ 64x64 5bit:
+
+1. Download from https://github.com/aravindsrinivas/flowpp/tree/master/flows_celeba.
+2. Convert .npy files to .pth using `preprocessing/convert_to_pth`.
+3. Place in `data/celebahq64_5bit`.
+
+CelebAHQ 256x256:
+```
+# Download Glow's preprocessed dataset.
+wget https://storage.googleapis.com/glow-demo/data/celeba-tfr.tar
+tar -C data/celebahq -xvf celeb-tfr.tar
+python extract_celeba_from_tfrecords
+```
+
+## Density Estimation Experiments
+
+MNIST:
+```
+python train_img.py --data mnist --imagesize 28 --actnorm True --wd 0 --save experiments/mnist
+```
+
+CIFAR10:
+```
+python train_img.py --data cifar10 --actnorm True --save experiments/cifar10
+```
+
+ImageNet 32x32:
+```
+python train_img.py --data imagenet32 --actnorm True --nblocks 32-32-32 --save experiments/imagenet32
+```
+
+ImageNet 64x64:
+```
+python train_img.py --data imagenet64 --imagesize 64 --actnorm True --nblocks 32-32-32 --factor-out True --squeeze-first True --save experiments/imagenet64
+```
+
+CelebAHQ 256x256:
+```
+python train_img.py --data celebahq --imagesize 256 --nbits 5 --actnorm True --act elu --batchsize 8 --update-freq 5 --n-exact-terms 8 --fc-end False --factor-out True --squeeze-first True --nblocks 16-16-16-16-16-16 --save experiments/celebahq256
+```
+
+## Pretrained Models
+
+Model checkpoints can be downloaded from [releases](https://github.com/rtqichen/residual-flows/releases/latest).
+
+Use the argument `--resume [checkpt.pth]` to evaluate or sample from the model. 
+
+Each checkpoint contains two sets of parameters, one from training and one containing the exponential moving average (EMA) accumulated over the course of training. Scripts will automatically use the EMA parameters for evaluation and sampling.
+
+## BibTeX
+```
+@inproceedings{chen2019residualflows,
+  title={Residual Flows for Invertible Generative Modeling},
+  author={Chen, Ricky T. Q. and Behrmann, Jens and Duvenaud, David and Jacobsen, J{\"{o}}rn{-}Henrik},
+  journal={CoRR},
+  volume={abs/1906.02735},
+  year={2019}
+}
+```

lib/datasets.py

@@ -0,0 +1,103 @@
+import torch
+import torchvision.datasets as vdsets
+
+
+class Dataset(object):
+
+    def __init__(self, loc, transform=None, in_mem=True):
+        self.in_mem = in_mem
+        self.dataset = torch.load(loc)
+        if in_mem: self.dataset = self.dataset.float().div(255)
+        self.transform = transform
+
+    def __len__(self):
+        return self.dataset.size(0)
+
+    @property
+    def ndim(self):
+        return self.dataset.size(1)
+
+    def __getitem__(self, index):
+        x = self.dataset[index]
+        if not self.in_mem: x = x.float().div(255)
+        x = self.transform(x) if self.transform is not None else x
+        return x, 0
+
+
+class MNIST(object):
+
+    def __init__(self, dataroot, train=True, transform=None):
+        self.mnist = vdsets.MNIST(dataroot, train=train, download=True, transform=transform)
+
+    def __len__(self):
+        return len(self.mnist)
+
+    @property
+    def ndim(self):
+        return 1
+
+    def __getitem__(self, index):
+        return self.mnist[index]
+
+
+class CIFAR10(object):
+
+    def __init__(self, dataroot, train=True, transform=None):
+        self.cifar10 = vdsets.CIFAR10(dataroot, train=train, download=True, transform=transform)
+
+    def __len__(self):
+        return len(self.cifar10)
+
+    @property
+    def ndim(self):
+        return 3
+
+    def __getitem__(self, index):
+        return self.cifar10[index]
+
+
+class CelebA5bit(object):
+
+    LOC = 'data/celebahq64_5bit/celeba_full_64x64_5bit.pth'
+
+    def __init__(self, train=True, transform=None):
+        self.dataset = torch.load(self.LOC).float().div(31)
+        if not train:
+            self.dataset = self.dataset[:5000]
+        self.transform = transform
+
+    def __len__(self):
+        return self.dataset.size(0)
+
+    @property
+    def ndim(self):
+        return self.dataset.size(1)
+
+    def __getitem__(self, index):
+        x = self.dataset[index]
+        x = self.transform(x) if self.transform is not None else x
+        return x, 0
+
+
+class CelebAHQ(Dataset):
+    TRAIN_LOC = 'data/celebahq/celeba256_train.pth'
+    TEST_LOC = 'data/celebahq/celeba256_validation.pth'
+
+    def __init__(self, train=True, transform=None):
+        return super(CelebAHQ, self).__init__(self.TRAIN_LOC if train else self.TEST_LOC, transform)
+
+
+class Imagenet32(Dataset):
+    TRAIN_LOC = 'data/imagenet32/train_32x32.pth'
+    TEST_LOC = 'data/imagenet32/valid_32x32.pth'
+
+    def __init__(self, train=True, transform=None):
+        return super(Imagenet32, self).__init__(self.TRAIN_LOC if train else self.TEST_LOC, transform)
+
+
+class Imagenet64(Dataset):
+    TRAIN_LOC = 'data/imagenet64/train_64x64.pth'
+    TEST_LOC = 'data/imagenet64/valid_64x64.pth'
+
+    def __init__(self, train=True, transform=None):
+        return super(Imagenet64, self).__init__(self.TRAIN_LOC if train else self.TEST_LOC, transform, in_mem=False)

lib/layers/__init__.py

@@ -0,0 +1,8 @@
+from .act_norm import *
+from .container import *
+from .coupling import *
+from .elemwise import *
+from .iresblock import *
+from .normalization import *
+from .squeeze import *
+from .glow import *

lib/layers/act_norm.py

@@ -0,0 +1,79 @@
+import torch
+import torch.nn as nn
+from torch.nn import Parameter
+
+__all__ = ['ActNorm1d', 'ActNorm2d']
+
+
+class ActNormNd(nn.Module):
+
+    def __init__(self, num_features, eps=1e-12):
+        super(ActNormNd, self).__init__()
+        self.num_features = num_features
+        self.eps = eps
+        self.weight = Parameter(torch.Tensor(num_features))
+        self.bias = Parameter(torch.Tensor(num_features))
+        self.register_buffer('initialized', torch.tensor(0))
+
+    @property
+    def shape(self):
+        raise NotImplementedError
+
+    def forward(self, x, logpx=None):
+        c = x.size(1)
+
+        if not self.initialized:
+            with torch.no_grad():
+                # compute batch statistics
+                x_t = x.transpose(0, 1).contiguous().view(c, -1)
+                batch_mean = torch.mean(x_t, dim=1)
+                batch_var = torch.var(x_t, dim=1)
+
+                # for numerical issues
+                batch_var = torch.max(batch_var, torch.tensor(0.2).to(batch_var))
+
+                self.bias.data.copy_(-batch_mean)
+                self.weight.data.copy_(-0.5 * torch.log(batch_var))
+                self.initialized.fill_(1)
+
+        bias = self.bias.view(*self.shape).expand_as(x)
+        weight = self.weight.view(*self.shape).expand_as(x)
+
+        y = (x + bias) * torch.exp(weight)
+
+        if logpx is None:
+            return y
+        else:
+            return y, logpx - self._logdetgrad(x)
+
+    def inverse(self, y, logpy=None):
+        assert self.initialized
+        bias = self.bias.view(*self.shape).expand_as(y)
+        weight = self.weight.view(*self.shape).expand_as(y)
+
+        x = y * torch.exp(-weight) - bias
+
+        if logpy is None:
+            return x
+        else:
+            return x, logpy + self._logdetgrad(x)
+
+    def _logdetgrad(self, x):
+        return self.weight.view(*self.shape).expand(*x.size()).contiguous().view(x.size(0), -1).sum(1, keepdim=True)
+
+    def __repr__(self):
+        return ('{name}({num_features})'.format(name=self.__class__.__name__, **self.__dict__))
+
+
+class ActNorm1d(ActNormNd):
+
+    @property
+    def shape(self):
+        return [1, -1]
+
+
+class ActNorm2d(ActNormNd):
+
+    @property
+    def shape(self):
+        return [1, -1, 1, 1]

lib/layers/base/__init__.py

@@ -0,0 +1,3 @@
+from .activations import *
+from .lipschitz import *
+from .mixed_lipschitz import *

lib/layers/base/activations.py

@@ -0,0 +1,75 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class Identity(nn.Module):
+
+    def forward(self, x):
+        return x
+
+
+class FullSort(nn.Module):
+
+    def forward(self, x):
+        return torch.sort(x, 1)[0]
+
+
+class MaxMin(nn.Module):
+
+    def forward(self, x):
+        b, d = x.shape
+        max_vals = torch.max(x.view(b, d // 2, 2), 2)[0]
+        min_vals = torch.min(x.view(b, d // 2, 2), 2)[0]
+        return torch.cat([max_vals, min_vals], 1)
+
+
+class LipschitzCube(nn.Module):
+
+    def forward(self, x):
+        return (x >= 1).to(x) * (x - 2 / 3) + (x <= -1).to(x) * (x + 2 / 3) + ((x > -1) * (x < 1)).to(x) * x**3 / 3
+
+
+class SwishFn(torch.autograd.Function):
+
+    @staticmethod
+    def forward(ctx, x, beta):
+        beta_sigm = torch.sigmoid(beta * x)
+        output = x * beta_sigm
+        ctx.save_for_backward(x, output, beta)
+        return output / 1.1
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        x, output, beta = ctx.saved_tensors
+        beta_sigm = output / x
+        grad_x = grad_output * (beta * output + beta_sigm * (1 - beta * output))
+        grad_beta = torch.sum(grad_output * (x * output - output * output)).expand_as(beta)
+        return grad_x / 1.1, grad_beta / 1.1
+
+
+class Swish(nn.Module):
+
+    def __init__(self):
+        super(Swish, self).__init__()
+        self.beta = nn.Parameter(torch.tensor([0.5]))
+
+    def forward(self, x):
+        return (x * torch.sigmoid_(x * F.softplus(self.beta))).div_(1.1)
+
+
+if __name__ == '__main__':
+
+    m = Swish()
+    xx = torch.linspace(-5, 5, 1000).requires_grad_(True)
+    yy = m(xx)
+    dd, dbeta = torch.autograd.grad(yy.sum() * 2, [xx, m.beta])
+
+    import matplotlib.pyplot as plt
+
+    plt.plot(xx.detach().numpy(), yy.detach().numpy(), label='Func')
+    plt.plot(xx.detach().numpy(), dd.detach().numpy(), label='Deriv')
+    plt.plot(xx.detach().numpy(), torch.max(dd.detach().abs() - 1, torch.zeros_like(dd)).numpy(), label='|Deriv| > 1')
+    plt.legend()
+    plt.tight_layout()
+    plt.show()