PicelShuffle1d added

Author: Zhangyanbo

src/INN/INN.py

@@ -9,6 +9,7 @@
 import INN.INNAbstract as INNAbstract
 import INN.cnn as cnn
 import torch.nn.functional as F
+import INN.pixel_shuffle_1d as ps
 
 iResNetModule = INNAbstract.iResNetModule
 
@@ -134,6 +135,24 @@ def PixelUnshuffle(self, x):
         return self.unshuffle(x)
 
 
+class PixelShuffle1d(INNAbstract.PixelShuffleModule):
+    '''
+    2d invertible pixel shuffle, using the built-in method
+    from pytorch. (nn.PixelShuffle, and nn.PixelUnshuffle)
+    '''
+    def __init__(self, r):
+        super(PixelShuffle1d, self).__init__()
+        self.r = r
+        self.shuffle = ps.PixelShuffle1D(r)
+        self.unshuffle = ps.PixelUnshuffle1D(r)
+    
+    def PixelShuffle(self, x):
+        return self.shuffle(x)
+    
+    def PixelUnshuffle(self, x):
+        return self.unshuffle(x)
+
+
 class BatchNorm1d(nn.BatchNorm1d, INNAbstract.INNModule):
     def __init__(self, dim, requires_grad=True):
         INNAbstract.INNModule.__init__(self)

src/INN/INNAbstract.py

@@ -131,14 +131,14 @@ def PixelShuffle(self, x):
     def PixelUnshuffle(self, x):
         pass
 
-    def forward(self, x, log_p0, log_det_J):
+    def forward(self, x, log_p0=0, log_det_J=0):
         # The log(p_0) and log|det J| will not change under this transformation
         if self.compute_p:
             return self.PixelUnshuffle(x), log_p0, log_det_J
         else:
             return self.PixelUnshuffle(x)
 
-    def inverse(self, y, num_iter=100):
+    def inverse(self, y, **args):
         return self.PixelShuffle(y)
 
 

src/INN/pixel_shuffle_1d.py

@@ -0,0 +1,51 @@
+import torch
+
+# source: https://github.com/serkansulun/pytorch-pixelshuffle1d/blob/master/pixelshuffle1d.py
+# "long" and "short" denote longer and shorter samples
+
+class PixelShuffle1D(torch.nn.Module):
+    """
+    1D pixel shuffler. https://arxiv.org/pdf/1609.05158.pdf
+    Upscales sample length, downscales channel length
+    "short" is input, "long" is output
+    """
+    def __init__(self, upscale_factor):
+        super(PixelShuffle1D, self).__init__()
+        self.upscale_factor = upscale_factor
+
+    def forward(self, x):
+        batch_size = x.shape[0]
+        short_channel_len = x.shape[1]
+        short_width = x.shape[2]
+
+        long_channel_len = short_channel_len // self.upscale_factor
+        long_width = self.upscale_factor * short_width
+
+        x = x.contiguous().view([batch_size, self.upscale_factor, long_channel_len, short_width])
+        x = x.permute(0, 2, 3, 1).contiguous()
+        x = x.view(batch_size, long_channel_len, long_width)
+
+        return x
+
+class PixelUnshuffle1D(torch.nn.Module):
+    """
+    Inverse of 1D pixel shuffler
+    Upscales channel length, downscales sample length
+    "long" is input, "short" is output
+    """
+    def __init__(self, downscale_factor):
+        super(PixelUnshuffle1D, self).__init__()
+        self.downscale_factor = downscale_factor
+
+    def forward(self, x):
+        batch_size = x.shape[0]
+        long_channel_len = x.shape[1]
+        long_width = x.shape[2]
+
+        short_channel_len = long_channel_len * self.downscale_factor
+        short_width = long_width // self.downscale_factor
+
+        x = x.contiguous().view([batch_size, long_channel_len, short_width, self.downscale_factor])
+        x = x.permute(0, 3, 1, 2).contiguous()
+        x = x.view([batch_size, short_channel_len, short_width])
+        return x