feat: initial for IDF

integer_discrete_flows/models/generative_flows.py

@@ -0,0 +1,175 @@
+"""
+Collection of flow strategies
+"""
+
+from __future__ import print_function
+
+import torch
+import numpy as np
+from models.utils import Base
+from .priors import SplitPrior
+from .coupling import Coupling
+
+
+UNIT_TESTING = False
+
+
+def space_to_depth(x):
+    xs = x.size()
+    # Pick off every second element
+    x = x.view(xs[0], xs[1], xs[2] // 2, 2, xs[3] // 2, 2)
+    # Transpose picked elements next to channels.
+    x = x.permute((0, 1, 3, 5, 2, 4)).contiguous()
+    # Combine with channels.
+    x = x.view(xs[0], xs[1] * 4, xs[2] // 2, xs[3] // 2)
+    return x
+
+
+def depth_to_space(x):
+    xs = x.size()
+    # Pick off elements from channels
+    x = x.view(xs[0], xs[1] // 4, 2, 2, xs[2], xs[3])
+    # Transpose picked elements next to HW dimensions.
+    x = x.permute((0, 1, 4, 2, 5, 3)).contiguous()
+    # Combine with HW dimensions.
+    x = x.view(xs[0], xs[1] // 4, xs[2] * 2, xs[3] * 2)
+    return x
+
+
+def int_shape(x):
+    return list(map(int, x.size()))
+
+
+class Flatten(Base):
+    def forward(self, x):
+        return x.view(x.size(0), -1)
+
+
+class Reshape(Base):
+    def __init__(self, shape):
+        super().__init__()
+        self.shape = shape
+
+    def forward(self, x):
+        return x.view(x.size(0), *self.shape)
+
+
+class Reverse(Base):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, z, reverse=False):
+        flip_idx = torch.arange(z.size(1) - 1, -1, -1).long()
+        z = z[:, flip_idx, :, :]
+        return z
+
+
+class Permute(Base):
+    def __init__(self, n_channels):
+        super().__init__()
+
+        permutation = np.arange(n_channels, dtype='int')
+        np.random.shuffle(permutation)
+
+        permutation_inv = np.zeros(n_channels, dtype='int')
+        permutation_inv[permutation] = np.arange(n_channels, dtype='int')
+
+        self.permutation = torch.from_numpy(permutation)
+        self.permutation_inv = torch.from_numpy(permutation_inv)
+
+    def forward(self, z, ldj, reverse=False):
+        if not reverse:
+            z = z[:, self.permutation, :, :]
+        else:
+            z = z[:, self.permutation_inv, :, :]
+
+        return z, ldj
+
+    def InversePermute(self):
+        inv_permute = Permute(len(self.permutation))
+        inv_permute.permutation = self.permutation_inv
+        inv_permute.permutation_inv = self.permutation
+        return inv_permute
+
+
+class Squeeze(Base):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, z, ldj, reverse=False):
+        if not reverse:
+            z = space_to_depth(z)
+        else:
+            z = depth_to_space(z)
+        return z, ldj
+
+
+class GenerativeFlow(Base):
+    def __init__(self, n_channels, height, width, args):
+        super().__init__()
+        layers = []
+        layers.append(Squeeze())
+        n_channels *= 4
+        height //= 2
+        width //= 2
+
+        for level in range(args.n_levels):
+
+            for i in range(args.n_flows):
+                perm_layer = Permute(n_channels)
+                layers.append(perm_layer)
+
+                layers.append(
+                    Coupling(n_channels, height, width, args))
+
+            if level < args.n_levels - 1:
+                if args.splitprior_type != 'none':
+                    # Standard splitprior
+                    factor_out = n_channels // 2
+                    layers.append(SplitPrior(n_channels, factor_out, height, width, args))
+                    n_channels = n_channels - factor_out
+
+                layers.append(Squeeze())
+                n_channels *= 4
+                height //= 2
+                width //= 2
+
+        self.layers = torch.nn.ModuleList(layers)
+        self.z_size = (n_channels, height, width)
+
+    def forward(self, z, ldj, pys=(), ys=(), reverse=False):
+        if not reverse:
+            for l, layer in enumerate(self.layers):
+                if isinstance(layer, (SplitPrior)):
+                    py, y, z, ldj = layer(z, ldj)
+                    pys += (py,)
+                    ys += (y,)
+
+                else:
+                    z, ldj = layer(z, ldj)
+
+        else:
+            for l, layer in reversed(list(enumerate(self.layers))):
+                if isinstance(layer, (SplitPrior)):
+                    if len(ys) > 0:
+                        z, ldj = layer.inverse(z, ldj, y=ys[-1])
+                        # Pop last element
+                        ys = ys[:-1]
+                    else:
+                        z, ldj = layer.inverse(z, ldj, y=None)
+
+                else:
+                    z, ldj = layer(z, ldj, reverse=True)
+
+        return z, ldj, pys, ys
+
+    def decode(self, z, ldj, state, decode_fn):
+
+        for l, layer in reversed(list(enumerate(self.layers))):
+            if isinstance(layer, SplitPrior):
+                z, ldj, state = layer.decode(z, ldj, state, decode_fn)
+
+            else:
+                z, ldj = layer(z, ldj, reverse=True)
+
+        return z, ldj