Dev vs (#15)

* Added function to create periodic features by concatenating sin and cos of features

* Removed unused referenced

* Added flow test module, test for glow

* Fixed mixing warning

* Old test module replaced

* Test for coupling layers added

* Added test for planar flow and nsf wrapper

* Fixed multidim planar flow

* Fix radial multidim case

* Added tests for residual flow

* Added link to docu in readme

* Compatibility with most recent pytorch version

README.md

@@ -8,7 +8,8 @@
 
 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). There are several sample use cases implemented in the 
+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),
 a [VAE](https://github.com/VincentStimper/normalizing-flows/blob/master/example/vae.py), and

normflows/flows/affine/autoregressive_test.py

@@ -2,12 +2,12 @@
 import unittest
 
 from normflows.flows.affine import autoregressive
-from normflows.flows.neural_spline.flow_test import FlowTest
+from normflows.flows.flow_test import FlowTest
 
 
 class MaskedAffineAutoregressiveTest(FlowTest):
-    def test_forward(self):
-        batch_size = 10
+    def test_maf(self):
+        batch_size = 3
         features = 20
         inputs = torch.randn(batch_size, features)
         for use_residual_blocks, random_mask in [
@@ -18,61 +18,14 @@ def test_forward(self):
             with self.subTest(
                 use_residual_blocks=use_residual_blocks, random_mask=random_mask
             ):
-                transform = autoregressive.MaskedAffineAutoregressive(
+                flow = autoregressive.MaskedAffineAutoregressive(
                     features=features,
                     hidden_features=30,
                     num_blocks=5,
                     use_residual_blocks=use_residual_blocks,
                     random_mask=random_mask,
                 )
-                outputs, logabsdet = transform(inputs)
-                self.assert_tensor_is_good(outputs, [batch_size, features])
-                self.assert_tensor_is_good(logabsdet, [batch_size])
-
-    def test_inverse(self):
-        batch_size = 10
-        features = 20
-        inputs = torch.randn(batch_size, features)
-        for use_residual_blocks, random_mask in [
-            (False, False),
-            (False, True),
-            (True, False),
-        ]:
-            with self.subTest(
-                use_residual_blocks=use_residual_blocks, random_mask=random_mask
-            ):
-                transform = autoregressive.MaskedAffineAutoregressive(
-                    features=features,
-                    hidden_features=30,
-                    num_blocks=5,
-                    use_residual_blocks=use_residual_blocks,
-                    random_mask=random_mask,
-                )
-                outputs, logabsdet = transform.inverse(inputs)
-                self.assert_tensor_is_good(outputs, [batch_size, features])
-                self.assert_tensor_is_good(logabsdet, [batch_size])
-
-    def test_forward_inverse_are_consistent(self):
-        batch_size = 10
-        features = 20
-        inputs = torch.randn(batch_size, features)
-        self.eps = 1e-6
-        for use_residual_blocks, random_mask in [
-            (False, False),
-            (False, True),
-            (True, False),
-        ]:
-            with self.subTest(
-                use_residual_blocks=use_residual_blocks, random_mask=random_mask
-            ):
-                transform = autoregressive.MaskedAffineAutoregressive(
-                    features=features,
-                    hidden_features=30,
-                    num_blocks=5,
-                    use_residual_blocks=use_residual_blocks,
-                    random_mask=random_mask,
-                )
-                self.assert_forward_inverse_are_consistent(transform, inputs)
+                self.checkForwardInverse(flow, inputs)
 
 
 if __name__ == "__main__":

normflows/flows/affine/coupling.py

@@ -2,7 +2,7 @@
 import torch
 from torch import nn
 
-from ..base import Flow
+from ..base import Flow, zero_log_det_like_z
 from ..reshape import Split, Merge
 
 
@@ -140,8 +140,8 @@ def forward(self, z):
             else:
                 raise NotImplementedError("This scale map is not implemented.")
         else:
-            z2 += param
-            log_det = 0
+            z2 = z2 + param
+            log_det = zero_log_det_like_z(z2)
         return [z1, z2], log_det
 
     def inverse(self, z):
@@ -164,8 +164,8 @@ def inverse(self, z):
             else:
                 raise NotImplementedError("This scale map is not implemented.")
         else:
-            z2 -= param
-            log_det = 0
+            z2 = z2 - param
+            log_det = zero_log_det_like_z(z2)
         return [z1, z2], log_det
 
 

normflows/flows/affine/coupling_test.py

@@ -0,0 +1,46 @@
+import unittest
+import torch
+
+from torch.testing import assert_close
+from normflows.flows import MaskedAffineFlow, CCAffineConst
+from normflows.nets import MLP
+from normflows.flows.flow_test import FlowTest
+
+
+class CouplingTest(FlowTest):
+    def test_mask_affine(self):
+        batch_size = 5
+        for latent_size in [2, 7]:
+            with self.subTest(latent_size=latent_size):
+                b = torch.Tensor([1 if i % 2 == 0 else 0 for i 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)
+                flow = MaskedAffineFlow(b, t, s)
+                inputs = torch.randn((batch_size, latent_size))
+                self.checkForwardInverse(flow, inputs)
+
+    def test_cc_affine(self):
+        batch_size = 5
+        for shape in [(5,), (2, 3, 4)]:
+            for num_classes in [2, 5]:
+                with self.subTest(shape=shape, num_classes=num_classes):
+                    flow = CCAffineConst(shape, num_classes)
+                    x = torch.randn((batch_size,) + shape)
+                    y = torch.rand((batch_size,) + (num_classes,))
+                    x_, log_det = flow(x, y)
+                    x__, log_det_ = flow(x_, y)
+
+                    assert x_.dtype == x.dtype
+                    assert x__.dtype == x.dtype
+
+                    assert x_.shape == x.shape
+                    assert x__.shape == x.shape
+
+                    assert_close(x__, x)
+                    id_ld = log_det + log_det_
+                    assert_close(id_ld, torch.zeros_like(id_ld))
+
+
+
+if __name__ == "__main__":
+    unittest.main()

normflows/flows/affine/glow.py

@@ -49,11 +49,11 @@ def __init__(
         kernel_size = (3, 1, 3)
         num_param = 2 if scale else 1
         if "channel" == split_mode:
-            channels_ = (channels // 2,) + 2 * (hidden_channels,)
-            channels_ += (num_param * ((channels + 1) // 2),)
-        elif "channel_inv" == split_mode:
             channels_ = ((channels + 1) // 2,) + 2 * (hidden_channels,)
             channels_ += (num_param * (channels // 2),)
+        elif "channel_inv" == split_mode:
+            channels_ = (channels // 2,) + 2 * (hidden_channels,)
+            channels_ += (num_param * ((channels + 1) // 2),)
         elif "checkerboard" in split_mode:
             channels_ = (channels,) + 2 * (hidden_channels,)
             channels_ += (num_param * channels,)

normflows/flows/affine/glow_test.py

@@ -0,0 +1,32 @@
+import unittest
+import torch
+
+from normflows.flows import GlowBlock
+from normflows.flows.flow_test import FlowTest
+
+
+class GlowTest(FlowTest):
+    def test_glow(self):
+        img_size = (4, 4)
+        hidden_channels = 8
+        for batch_size, channels, scale, split_mode, use_lu, net_actnorm in [
+            (1, 3, True, "channel", True, False),
+            (2, 3, True, "channel_inv", True, False),
+            (1, 4, True, "channel_inv", True, True),
+            (2, 4, True, "channel", True, False),
+            (1, 4, False, "channel", False, False),
+            (1, 4, True, "checkerboard", True, True),
+            (3, 5, False, "checkerboard", False, True)
+        ]:
+            with self.subTest(batch_size=batch_size, channels=channels,
+                              scale=scale, split_mode=split_mode,
+                              use_lu=use_lu, net_actnorm=net_actnorm):
+                inputs = torch.rand((batch_size, channels) + img_size)
+                flow = GlowBlock(channels, hidden_channels,
+                                 scale=scale, split_mode=split_mode,
+                                 use_lu=use_lu, net_actnorm=net_actnorm)
+                self.checkForwardInverse(flow, inputs)
+
+
+if __name__ == "__main__":
+    unittest.main()

normflows/flows/base.py

@@ -76,3 +76,7 @@ def forward(self, inputs):
     def inverse(self, inputs):
         funcs = (flow.inverse for flow in self._flows[::-1])
         return self._cascade(inputs, funcs)
+
+
+def zero_log_det_like_z(z):
+    return torch.zeros(z.shape[0], dtype=z.dtype, device=z.device)

normflows/flows/flow_test.py

@@ -0,0 +1,42 @@
+import unittest
+import torch
+
+from torch.testing import assert_close
+
+
+class FlowTest(unittest.TestCase):
+    """
+    Generic test case for flow modules
+    """
+    def assertClose(self, actual, expected, atol=None, rtol=None):
+        assert_close(actual, expected, atol=atol, rtol=rtol)
+
+    def checkForward(self, flow, inputs):
+        # Do forward transform
+        outputs, log_det = flow(inputs)
+        # Check type
+        assert outputs.dtype == inputs.dtype
+        # Check shape
+        assert outputs.shape == inputs.shape
+        # Return results
+        return outputs, log_det
+
+    def checkInverse(self, flow, inputs):
+        # Do inverse transform
+        outputs, log_det = flow.inverse(inputs)
+        # Check type
+        assert outputs.dtype == inputs.dtype
+        # Check shape
+        assert outputs.shape == inputs.shape
+        # Return results
+        return outputs, log_det
+
+    def checkForwardInverse(self, flow, inputs, atol=None, rtol=None):
+        # Check forward
+        outputs, log_det = self.checkForward(flow, inputs)
+        # Check inverse
+        input_, log_det_ = self.checkInverse(flow, outputs)
+        # Check identity
+        self.assertClose(input_, inputs, atol, rtol)
+        ld_id = log_det + log_det_
+        self.assertClose(ld_id, torch.zeros_like(ld_id), atol, rtol)

normflows/flows/mixing.py

@@ -70,7 +70,7 @@ def __init__(self, num_channels, use_lu=False):
         super().__init__()
         self.num_channels = num_channels
         self.use_lu = use_lu
-        Q = torch.qr(torch.randn(self.num_channels, self.num_channels))[0]
+        Q, _ = torch.linalg.qr(torch.randn(self.num_channels, self.num_channels))
         if use_lu:
             P, L, U = torch.lu_unpack(*Q.lu())
             self.register_buffer("P", P)  # remains fixed during optimization
@@ -149,7 +149,7 @@ def __init__(self, num_channels, use_lu=True):
         super().__init__()
         self.num_channels = num_channels
         self.use_lu = use_lu
-        Q = torch.qr(torch.randn(self.num_channels, self.num_channels))[0]
+        Q, _ = torch.linalg.qr(torch.randn(self.num_channels, self.num_channels))
         if use_lu:
             P, L, U = torch.lu_unpack(*Q.lu())
             self.register_buffer("P", P)  # remains fixed during optimization

normflows/flows/neural_spline/autoregressive.py

@@ -11,7 +11,7 @@
 from ..affine.autoregressive import Autoregressive
 from normflows.nets import made as made_module
 from normflows.utils import splines
-from normflows.utils.nn import PeriodicFeatures
+from normflows.utils.nn import PeriodicFeaturesElementwise
 
 
 class MaskedPiecewiseRationalQuadraticAutoregressive(Autoregressive):
@@ -50,7 +50,7 @@ def __init__(
                 scale_pf = np.pi / tail_bound[ind_circ]
             else:
                 scale_pf = np.pi / tail_bound
-            preprocessing = PeriodicFeatures(features, ind_circ, scale_pf)
+            preprocessing = PeriodicFeaturesElementwise(features, ind_circ, scale_pf)
         else:
             preprocessing = None
 

normflows/flows/neural_spline/autoregressive_test.py

@@ -1,31 +1,30 @@
 """
 Tests for the autoregressive transforms.
-Code taken from https://github.com/bayesiains/nsf
+Code partially taken from https://github.com/bayesiains/nsf
 """
 
 import torch
 import unittest
 
 from normflows.flows.neural_spline import autoregressive
-from normflows.flows.neural_spline.flow_test import FlowTest
+from normflows.flows.flow_test import FlowTest
 
 
 class MaskedPiecewiseRationalQuadraticAutoregressiveFlowTest(FlowTest):
-    def test_forward_inverse_are_consistent(self):
-        batch_size = 10
-        features = 20
+    def test_mprqas(self):
+        batch_size = 5
+        features = 10
         inputs = torch.rand(batch_size, features)
-        self.eps = 1e-3
 
-        transform = autoregressive.MaskedPiecewiseRationalQuadraticAutoregressive(
+        flow = autoregressive.MaskedPiecewiseRationalQuadraticAutoregressive(
             num_bins=10,
             features=features,
             hidden_features=30,
             num_blocks=5,
             use_residual_blocks=True,
         )
 
-        self.assert_forward_inverse_are_consistent(transform, inputs)
+        self.checkForwardInverse(flow, inputs)
 
 
 if __name__ == "__main__":

normflows/flows/neural_spline/coupling.py

@@ -7,7 +7,6 @@
 import numpy as np
 import torch
 from torch import nn
-from torch.nn import functional as F
 
 from ..base import Flow
 from ... import utils