Resolved merge conflict

README.md

@@ -1,4 +1,4 @@
-# Normalizing Flows
+# `normflows`: A PyTorch Package for Normalizing Flows
 
 [![documentation](https://github.com/VincentStimper/normalizing-flows/actions/workflows/mkdocs.yaml/badge.svg)](https://vincentstimper.github.io/normalizing-flows/)
 ![unit-tests](https://github.com/VincentStimper/normalizing-flows/actions/workflows/pytest.yaml/badge.svg)
@@ -11,7 +11,7 @@
 
 This is a PyTorch implementation of normalizing flows. Many popular flow architectures are implemented,
 see the [list below](#implemented-flows). The package can be easily [installed via pip](#installation).
-The basic usage is described [here](#usage), and a [full documentation](https://vincentstimper.github.io/normalizing-flows/). 
+The basic usage is described [here](#usage), and a [full documentation](https://vincentstimper.github.io/normalizing-flows/) 
 is available as well. There are several sample use cases implemented in the 
 [`example` folder](https://github.com/VincentStimper/normalizing-flows/tree/master/example), 
 including [Glow](https://github.com/VincentStimper/normalizing-flows/blob/master/example/glow.ipynb),
@@ -21,19 +21,21 @@ a [Residual Flow](https://github.com/VincentStimper/normalizing-flows/blob/maste
 
 ## Implemented Flows
 
-* Planar Flow ([Rezende & Mohamed, 2015](https://arxiv.org/abs/1505.05770))
-* Radial Flow ([Rezende & Mohamed, 2015](https://arxiv.org/abs/1505.05770))
-* NICE ([Dinh et al., 2014](https://arxiv.org/abs/1410.8516))
-* Real NVP ([Dinh et al., 2016](https://arxiv.org/abs/1605.08803))
-* Glow ([Kingma & Dhariwal, 2018](https://arxiv.org/abs/1807.03039))
-* Masked Autoregressive Flow ([Papamakarios et al., 2017](https://proceedings.neurips.cc/paper/2017/hash/6c1da886822c67822bcf3679d04369fa-Abstract.html))
-* Neural Spline Flow ([Durkan et al., 2019](https://arxiv.org/abs/1906.04032))
-* Circular Neural Spline Flow ([Rezende et al., 2020](http://proceedings.mlr.press/v119/rezende20a.html))
-* Residual Flow ([Chen et al., 2019](https://arxiv.org/abs/1906.02735))
-* Stochastic Normalizing Flows ([Wu et al., 2020](https://arxiv.org/abs/2002.06707))
+| Architecture | Reference                                                                                                                 |
+|--------------|---------------------------------------------------------------------------------------------------------------------------|
+| Planar Flow  | [Rezende & Mohamed, 2015](http://proceedings.mlr.press/v37/rezende15.html)                                                |
+| Radial Flow  | [Rezende & Mohamed, 2015](http://proceedings.mlr.press/v37/rezende15.html)                                                |
+| NICE         | [Dinh et al., 2014](https://arxiv.org/abs/1410.8516)                                                                      |
+| Real NVP     | [Dinh et al., 2017](https://openreview.net/forum?id=HkpbnH9lx)                                                            |
+| Glow         | [Kingma et al., 2018](https://proceedings.neurips.cc/paper/2018/hash/d139db6a236200b21cc7f752979132d0-Abstract.html)                                                                   |
+| Masked Autoregressive Flow | [Papamakarios et al., 2017](https://proceedings.neurips.cc/paper/2017/hash/6c1da886822c67822bcf3679d04369fa-Abstract.html) |
+| Neural Spline Flow | [Durkan et al., 2019](https://proceedings.neurips.cc/paper/2019/hash/7ac71d433f282034e088473244df8c02-Abstract.html)                                                                    |
+| Circular Neural Spline Flow | [Rezende et al., 2020](http://proceedings.mlr.press/v119/rezende20a.html)                                                 |
+| Residual Flow | [Chen et al., 2019](https://proceedings.neurips.cc/paper/2019/hash/5d0d5594d24f0f955548f0fc0ff83d10-Abstract.html)                                                                     |
+| Stochastic Normalizing Flow | [Wu et al., 2020](https://proceedings.neurips.cc/paper/2020/hash/41d80bfc327ef980528426fc810a6d7a-Abstract.html)                                                                       |
 
 Note that Neural Spline Flows with circular and non-circular coordinates
-are also supported.
+are supported as well.
 
 ## Installation
 

normflows/core.py

@@ -24,6 +24,66 @@ def __init__(self, q0, flows, p=None):
         self.flows = nn.ModuleList(flows)
         self.p = p
 
+    def forward(self, z):
+        """Transforms latent variable z to the flow variable x
+
+        Args:
+          z: Batch in the latent space
+
+        Returns:
+          Batch in the space of the target distribution
+        """
+        for flow in self.flows:
+            z, _ = flow(z)
+        return z
+
+    def forward_and_log_det(self, z):
+        """Transforms latent variable z to the flow variable x and
+        computes log determinant of the Jacobian
+
+        Args:
+          z: Batch in the latent space
+
+        Returns:
+          Batch in the space of the target distribution,
+          log determinant of the Jacobian
+        """
+        log_det = torch.zeros(len(z), device=z.device)
+        for flow in self.flows:
+            z, log_d = flow(z)
+            log_det -= log_d
+        return z, log_det
+
+    def inverse(self, x):
+        """Transforms flow variable x to the latent variable z
+
+        Args:
+          x: Batch in the space of the target distribution
+
+        Returns:
+          Batch in the latent space
+        """
+        for i in range(len(self.flows) - 1, -1, -1):
+            x, _ = self.flows[i].inverse(x)
+        return x
+
+    def inverse_and_log_det(self, x):
+        """Transforms flow variable x to the latent variable z and
+        computes log determinant of the Jacobian
+
+        Args:
+          x: Batch in the space of the target distribution
+
+        Returns:
+          Batch in the latent space, log determinant of the
+          Jacobian
+        """
+        log_det = torch.zeros(len(x), device=x.device)
+        for i in range(len(self.flows) - 1, -1, -1):
+            x, log_d = self.flows[i].inverse(x)
+            log_det += log_d
+        return x, log_det
+
     def forward_kld(self, x):
         """Estimates forward KL divergence, see [arXiv 1912.02762](https://arxiv.org/abs/1912.02762)
 
@@ -267,6 +327,7 @@ def forward_kld(self, x, y=None):
 
         Args:
           x: Batch sampled from target distribution
+          y: Batch of targets, if applicable
 
         Returns:
           Estimate of forward KL divergence averaged over batch
@@ -279,6 +340,9 @@ def forward(self, x, y=None):
         Args:
           x: Batch of data
           y: Batch of targets, if applicable
+
+        Returns:
+            Negative log-likelihood of the batch
         """
         return -self.log_prob(x, y)
 

normflows/core_test.py

@@ -0,0 +1,171 @@
+import unittest
+import torch
+
+from torch.testing import assert_close
+from normflows import NormalizingFlow, ClassCondFlow, \
+    MultiscaleFlow, NormalizingFlowVAE
+from normflows.flows import MaskedAffineFlow, GlowBlock, \
+    Merge, Squeeze
+from normflows.nets import MLP
+from normflows.distributions.base import DiagGaussian, \
+    ClassCondDiagGaussian
+from normflows.distributions.target import CircularGaussianMixture
+from normflows.distributions.encoder import NNDiagGaussian
+from normflows.distributions.decoder import NNDiagGaussianDecoder
+
+
+class CoreTest(unittest.TestCase):
+    def test_normalizing_flow(self):
+        batch_size = 5
+        latent_size = 2
+        for n_layers in [2, 5]:
+            with self.subTest(n_layers=n_layers):
+                # Construct flow model
+                layers = []
+                for i in range(n_layers):
+                    b = torch.Tensor([j if i % 2 == j % 2 else 0 for j in range(latent_size)])
+                    s = MLP([latent_size, 2 * latent_size, latent_size], init_zeros=True)
+                    t = MLP([latent_size, 2 * latent_size, latent_size], init_zeros=True)
+                    layers.append(MaskedAffineFlow(b, t, s))
+                base = DiagGaussian(latent_size)
+                target = CircularGaussianMixture()
+                model = NormalizingFlow(base, layers, target)
+                # Test log prob and sampling
+                inputs = torch.randn((batch_size, latent_size))
+                log_q = model.log_prob(inputs)
+                assert log_q.shape == (batch_size,)
+                s, log_q = model.sample(batch_size)
+                assert log_q.shape == (batch_size,)
+                assert s.shape == (batch_size, latent_size)
+                # Test losses
+                loss = model.forward_kld(inputs)
+                assert loss.dim() == 0
+                loss = model.reverse_kld(batch_size)
+                assert loss.dim() == 0
+                loss = model.reverse_kld(batch_size, score_fn=False)
+                assert loss.dim() == 0
+                loss = model.reverse_alpha_div(batch_size)
+                assert loss.dim() == 0
+                loss = model.reverse_alpha_div(batch_size, dreg=True)
+                assert loss.dim() == 0
+                # Test forward and inverse
+                outputs = model.forward(inputs)
+                inputs_ = model.inverse(outputs)
+                assert_close(inputs_, inputs)
+                outputs, log_det = model.forward_and_log_det(inputs)
+                inputs_, log_det_ = model.inverse_and_log_det(outputs)
+                assert_close(inputs_, inputs)
+                assert_close(log_det, -log_det_)
+
+    def test_cc_normalizing_flow(self):
+        batch_size = 5
+        latent_size = 2
+        n_layers = 2
+        n_classes = 3
+
+        # Construct flow model
+        layers = []
+        for i in range(n_layers):
+            b = torch.Tensor([j if i % 2 == j % 2 else 0 for j in range(latent_size)])
+            s = MLP([latent_size, 2 * latent_size, latent_size], init_zeros=True)
+            t = MLP([latent_size, 2 * latent_size, latent_size], init_zeros=True)
+            layers.append(MaskedAffineFlow(b, t, s))
+        base = ClassCondDiagGaussian(latent_size, n_classes)
+        model = ClassCondFlow(base, layers)
+
+        # Test model
+        x = torch.randn((batch_size, latent_size))
+        y = torch.randint(n_classes, (batch_size,))
+        log_q = model.log_prob(x, y)
+        assert log_q.shape == (batch_size,)
+        s, log_q = model.sample(x, y)
+        assert log_q.shape == (batch_size,)
+        assert s.shape == (batch_size, latent_size)
+        loss = model.forward_kld(x, y)
+        assert loss.dim() == 0
+
+    def test_multiscale_flow(self):
+        # Set parameters
+        batch_size = 5
+        K = 2
+        hidden_channels = 32
+        split_mode = 'channel'
+        scale = True
+        params = [(2, 1, (3, 8, 8)), (3, 5, (1, 16, 16))]
+
+        for L, num_classes, input_shape in params:
+            with self.subTest(L=L, num_classes=num_classes,
+                              input_shape=input_shape):
+                # Set up flows, distributions and merge operations
+                base = []
+                merges = []
+                flows = []
+                for i in range(L):
+                    flows_ = []
+                    for j in range(K):
+                        flows_ += [GlowBlock(input_shape[0] * 2 ** (L + 1 - i), hidden_channels,
+                                             split_mode=split_mode, scale=scale)]
+                    flows_ += [Squeeze()]
+                    flows += [flows_]
+                    if i > 0:
+                        merges += [Merge()]
+                        latent_shape = (input_shape[0] * 2 ** (L - i), input_shape[1] // 2 ** (L - i),
+                                        input_shape[2] // 2 ** (L - i))
+                    else:
+                        latent_shape = (input_shape[0] * 2 ** (L + 1), input_shape[1] // 2 ** L,
+                                        input_shape[2] // 2 ** L)
+                    base += [ClassCondDiagGaussian(latent_shape, num_classes)]
+
+                # Construct flow model
+                model = MultiscaleFlow(base, flows, merges)
+                # Test model
+                y = torch.randint(num_classes, (batch_size,))
+                x, log_q = model.sample(batch_size, y)
+                log_q_ = model.log_prob(x, y)
+                assert x.shape == (batch_size,) + (input_shape)
+                assert log_q.shape == (batch_size,)
+                assert log_q_.shape == (batch_size,)
+                assert log_q.dtype == x.dtype
+                assert log_q_.dtype == x.dtype
+                assert_close(log_q, log_q_)
+                fwd = model.forward(x, y)
+                fwd_kld = model.forward_kld(x, y)
+                assert_close(torch.mean(fwd), fwd_kld)
+
+
+    def test_normalizing_flow_vae(self):
+        batch_size = 5
+        n_dim = 10
+        n_layers = 2
+        n_bottleneck = 3
+        n_hidden_untis = 16
+        hidden_units_encoder = [n_dim, n_hidden_untis, n_bottleneck * 2]
+        hidden_units_decoder = [n_bottleneck, n_hidden_untis, 2 * n_dim]
+
+        # Construct flow model
+        layers = []
+        for i in range(n_layers):
+            b = torch.Tensor([j if i % 2 == j % 2 else 0 for j in range(n_bottleneck)])
+            s = MLP([n_bottleneck, 2 * n_bottleneck, n_bottleneck], init_zeros=True)
+            t = MLP([n_bottleneck, 2 * n_bottleneck, n_bottleneck], init_zeros=True)
+            layers.append(MaskedAffineFlow(b, t, s))
+        prior = torch.distributions.MultivariateNormal(torch.zeros(n_bottleneck),
+                                                       torch.eye(n_bottleneck))
+        encoder_nn = MLP(hidden_units_encoder)
+        encoder = NNDiagGaussian(encoder_nn)
+        decoder_nn = MLP(hidden_units_decoder)
+        decoder = NNDiagGaussianDecoder(decoder_nn)
+        model = NormalizingFlowVAE(prior, encoder, layers, decoder)
+
+        # Test model
+        for num_samples in [1, 4]:
+            with self.subTest(num_samples=num_samples):
+                x = torch.randn((batch_size, n_dim))
+                z, log_p, log_q = model(x, num_samples=num_samples)
+                assert z.shape == (batch_size, num_samples, n_bottleneck)
+                assert log_p.shape == (batch_size, num_samples)
+                assert log_q.shape == (batch_size, num_samples)
+
+
+if __name__ == "__main__":
+    unittest.main()

normflows/distributions/base.py

@@ -36,6 +36,18 @@ def log_prob(self, z):
         """
         raise NotImplementedError
 
+    def sample(self, num_samples=1, **kwargs):
+        """Samples from base distribution
+
+        Args:
+          num_samples: Number of samples to draw from the distriubtion
+
+        Returns:
+          Samples drawn from the distribution
+        """
+        z, _ = self.forward(num_samples, **kwargs)
+        return z
+
 
 class DiagGaussian(BaseDistribution):
     """
@@ -51,6 +63,8 @@ def __init__(self, shape, trainable=True):
         super().__init__()
         if isinstance(shape, int):
             shape = (shape,)
+        if isinstance(shape, list):
+            shape = tuple(shape)
         self.shape = shape
         self.n_dim = len(shape)
         self.d = np.prod(shape)
@@ -104,6 +118,8 @@ def __init__(self, ndim, ind, scale=None):
         """
         super().__init__()
         self.ndim = ndim
+        if isinstance(ind, int):
+            ind = [ind]
 
         # Set up indices and permutations
         self.ndim = ndim
@@ -176,6 +192,8 @@ def __init__(self, shape, num_classes):
         super().__init__()
         if isinstance(shape, int):
             shape = (shape,)
+        if isinstance(shape, list):
+            shape = tuple(shape)
         self.shape = shape
         self.n_dim = len(shape)
         self.perm = [self.n_dim] + list(range(self.n_dim))
@@ -251,6 +269,8 @@ def __init__(self, shape, num_classes=None, logscale_factor=3.0):
         # Save shape and related statistics
         if isinstance(shape, int):
             shape = (shape,)
+        if isinstance(shape, list):
+            shape = tuple(shape)
         self.shape = shape
         self.n_dim = len(shape)
         self.num_pix = np.prod(shape[1:])
@@ -373,6 +393,10 @@ def __init__(self, shape, affine_shape, num_classes=None):
           num_classes: Number of classes if the base is class conditional, None otherwise
         """
         super().__init__()
+        if isinstance(shape, int):
+            shape = (shape,)
+        if isinstance(shape, list):
+            shape = tuple(shape)
         self.shape = shape
         self.n_dim = len(shape)
         self.d = np.prod(shape)
@@ -421,6 +445,9 @@ def forward(self, num_samples=1, y=None):
         else:
             z, log_det = self.transform(z)
         log_p -= log_det
+        print(y)
+        print(eps)
+        print(z)
         return z, log_p
 
     def log_prob(self, z, y=None):
@@ -450,6 +477,8 @@ def log_prob(self, z, y=None):
             - 0.5 * self.d * np.log(2 * np.pi)
             - 0.5 * torch.sum(torch.pow(z, 2), dim=self.sum_dim)
         )
+        print(y)
+        print(z)
         return log_p