add ResizeFeatures function

Author: Zhangyanbo

setup.py

@@ -4,8 +4,8 @@
     long_description = fh.read()
 
 setuptools.setup(
-    name="INNLab", # Replace with your own username
-    version="0.0.2",
+    name="INNLab",
+    version="0.1.0",
     author="Yanbo Zhang",
     author_email="zhangybspm@gmail.com",
     description="A package for invertible neural networks",

src/INN/INN.py

@@ -159,9 +159,10 @@ def PixelUnshuffle(self, x):
 
 
 class BatchNorm1d(nn.BatchNorm1d, INNAbstract.INNModule):
-    def __init__(self, dim):
+    def __init__(self, dim, requires_grad=False):
         INNAbstract.INNModule.__init__(self)
         nn.BatchNorm1d.__init__(self, num_features=dim, affine=False)
+        self.requires_grad = requires_grad
 
     def forward(self, x, log_p=0, log_det_J=0):
 
@@ -171,7 +172,9 @@ def forward(self, x, log_p=0, log_det_J=0):
                 var = self.running_var # [dim]
             else:
                 # if in training
-                var = torch.var(x, dim=0, unbiased=False).detach() # [dim]
+                var = torch.var(x, dim=0, unbiased=False)#.detach() # [dim]
+                if not self.requires_grad:
+                    var = var.detach()
 
             x = super(BatchNorm1d, self).forward(x)
 
@@ -211,11 +214,11 @@ def inverse(self, y, **args):
 
 class RealNVP(INNAbstract.INNModule):
 
-    def __init__(self, dim=None, f_log_s=None, f_t=None, k=4, mask=None, clip=1):
+    def __init__(self, dim=None, f_log_s=None, f_t=None, k=4, mask=None, clip=1, activation_fn=None):
         super(RealNVP, self).__init__()
         if (f_log_s is None) and (f_t is None):
-            log_s = utilities.default_net(dim, k)#self.default_net(dim, k)
-            t = utilities.default_net(dim, k)#self.default_net(dim, k)
+            log_s = utilities.default_net(dim, k, activation_fn)#self.default_net(dim, k)
+            t = utilities.default_net(dim, k, activation_fn)#self.default_net(dim, k)
             self.net = utilities.combined_real_nvp(dim, log_s, t, mask, clip)
         else:
             self.net = utilities.combined_real_nvp(dim, f_log_s, f_t, mask, clip)
@@ -234,11 +237,11 @@ def inverse(self, y, **args):
 
 class NICE(INNAbstract.INNModule):
 
-    def __init__(self, dim=None, m=None, mask=None, k=4):
+    def __init__(self, dim=None, m=None, mask=None, k=4, activation_fn=None):
         super(NICE, self).__init__()
 
         if m is None:
-            m_ = utilities.default_net(dim, k)
+            m_ = utilities.default_net(dim, k, activation_fn)
             self.net = utilities.NICE(dim, m=m_, mask=mask)
         else:
             self.net = utilities.NICE(dim, m=m, mask=mask)
@@ -268,18 +271,18 @@ class Nonlinear(INNAbstract.INNModule):
     '''
     Nonlinear invertible block
     '''
-    def __init__(self, dim, method='NICE', m=None, mask=None, k=4, **args):
+    def __init__(self, dim, method='RealNVP', m=None, mask=None, k=4, activation_fn=None, **args):
         super(Nonlinear, self).__init__()
 
         self.method = method
         if method == 'NICE':
-            self.block = NICE(dim, m=m, mask=mask, k=k)
+            self.block = NICE(dim, m=m, mask=mask, k=k, activation_fn=activation_fn)
         if method == 'RealNVP':
             clip = _default_dict('clip', args, 1)
             f_log_s = _default_dict('f_log_s', args, None)
             f_t = _default_dict('f_t', args, None)
 
-            self.block = RealNVP(dim=dim, f_log_s=f_log_s, f_t=f_t, k=k, mask=mask, clip=clip)
+            self.block = RealNVP(dim=dim, f_log_s=f_log_s, f_t=f_t, k=k, mask=mask, clip=clip, activation_fn=activation_fn)
         if method == 'iResNet':
             g = _default_dict('g', args, None)
             beta = _default_dict('beta', args, 0.8)
@@ -293,4 +296,70 @@ def forward(self, x, log_p0=0, log_det_J=0):
         return self.block(x, log_p0, log_det_J)
 
     def inverse(self, y, **args):
-        return self.block.inverse(y, **args)
+        return self.block.inverse(y, **args)
+
+class ResizeFeatures(INNAbstract.INNModule):
+    '''
+    Resize for n-d input, include linear or multi-channel inputs
+    '''
+    def __init__(self, feature_in, feature_out, dist='normal'):
+        super(ResizeFeatures, self).__init__()
+        self.feature_in = feature_in
+        self.feature_out = feature_out
+
+        if dist == 'normal':
+            self.dist = utilities.NormalDistribution()
+        elif isinstance(dist, INNAbstract.Distribution):
+            self.dist = dist
+    
+    def resize(self, x, feature_in, feature_out):
+        '''
+        x has two kinds of shapes:
+            1. [feature_in]
+            2. [batch_size, feature_in, *]
+        '''
+        if len(x.shape) == 1:
+            # [feature_in]
+            if x.shape[0] != self.feature_in:
+                raise Exception(f'Expect to get {self.feature_in} features, but got {x.shape[0]}.')
+            y, z = x[:feature_out], x[feature_out:]
+        
+        if len(x.shape) >= 2:
+            # [batch_size, feature_in, *]
+            if x.shape[1] != self.feature_in:
+                raise Exception(f'Expect to get {self.feature_in} features, but got {x.shape[1]}.')
+            y, z = x[:, :feature_out], x[:, feature_out:]
+        
+        return y, z
+
+    def forward(self, x, log_p0=0, log_det_J=0):
+        x, z = self.resize(x, self.feature_in, self.feature_out)
+        if self.compute_p:
+            p = self.dist.logp(z)
+            return x, log_p0 + p, log_det_J
+        else:
+            return x
+    
+    def inverse(self, y, **args):
+        '''
+        y has two kinds of shapes:
+            1. [feature_in]
+            2. [batch_size, feature_in, *]
+        '''
+        if len(y.shape) == 1:
+            # [feature_in]
+            if y.shape[0] != self.feature_out:
+                raise Exception(f'Expect to get {self.feature_out} features, but got {y.shape[0]}.')
+            z = self.dist.sample(self.feature_in-self.feature_out).to(y.device)
+            y = torch.cat([y, z])
+        
+        if len(y.shape) >= 2:
+            # [batch_size, feature_in, *]
+            if y.shape[1] != self.feature_out:
+                raise Exception(f'Expect to get {self.feature_out} features, but got {y.shape[1]}.')
+            shape = list(y.shape)
+            shape[1] = self.feature_in-self.feature_out
+            z = self.dist.sample(shape).to(y.device)
+            y = torch.cat([y, z], dim=1)
+        
+        return y

src/INN/INNAbstract.py

@@ -140,4 +140,21 @@ def forward(self, x, log_p0, log_det_J):
             return self.PixelUnshuffle(x)
 
     def inverse(self, y, num_iter=100):
-        return self.PixelShuffle(y)
+        return self.PixelShuffle(y)
+
+
+class Distribution(nn.Module):
+
+    def __init__(self):
+        super(Distribution, self).__init__()
+    
+    def logp(self, x):
+        raise NotImplementedError('logp() not implemented')
+    
+    def sample(self, shape):
+        raise NotImplementedError('sample() not implemented')
+
+    def forward(self, x):
+        x = self.logp(x)
+        
+        return x

src/INN/utilities.py

@@ -138,7 +138,7 @@ def forward(self, x):
         return x
 
 
-class NormalDistribution(nn.Module):
+class NormalDistribution(INNAbstract.Distribution):
     '''
     Generate normal distribution and compute log probablity
     '''
@@ -166,10 +166,6 @@ def logp(self, x):
     def sample(self, shape):
         return torch.randn(shape)
 
-    def forward(self, x):
-        x = self.logp(x)
-        
-        return x
 
 def permutation_matrix(dim):
     # generate a permuation matrix
@@ -396,21 +392,38 @@ def inverse(self, y):
 
 
 class default_net(nn.Module):
-    def __init__(self, dim, k):
+    def __init__(self, dim, k, activation_fn=None):
         super(default_net, self).__init__()
-        self.net = self.default_net(dim, k)
+        self.activation_fn = activation_fn
+        self.net = self.default_net(dim, k, activation_fn)
 
-    def default_net(self, dim, k):
-        block = nn.Sequential(nn.Linear(dim, k * dim), nn.LeakyReLU(),
-                              nn.Linear(k * dim, k * dim), nn.LeakyReLU(),
+    def default_net(self, dim, k, activation_fn):
+        if activation_fn == None:
+            ac = nn.LeakyReLU
+        else:
+            ac = activation_fn
+        
+        block = nn.Sequential(nn.Linear(dim, k * dim), ac(),
+                              nn.Linear(k * dim, k * dim), ac(),
                               nn.Linear(k * dim, dim))
         block.apply(self.init_weights)
         return block
 
     def init_weights(self, m):
+        nonlinearity = 'leaky_relu' # set to leaky_relu by default
+
+        if self.activation_fn is nn.ReLU:
+            nonlinearity = 'leaky_relu'
+        if self.activation_fn is nn.SELU:
+            nonlinearity = 'selu'
+        if self.activation_fn is nn.Tanh:
+            nonlinearity = 'tanh'
+        if self.activation_fn is nn.Sigmoid:
+            nonlinearity = 'sigmoid'
+        
         if type(m) == nn.Linear:
             # doing Kaiming initialization
-            torch.nn.init.kaiming_normal_(m.weight.data, nonlinearity='leaky_relu')
+            torch.nn.init.kaiming_normal_(m.weight.data, nonlinearity=nonlinearity)
             torch.nn.init.zeros_(m.bias.data)
 
     def forward(self, x):

tests/quick_tests.ipynb

@@ -461,6 +461,96 @@
     "bn(x)"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": 24,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2021-04-27T00:31:48.829977Z",
+     "start_time": "2021-04-27T00:31:48.826448Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "x = torch.randn((3,3,3))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 29,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2021-04-27T00:32:17.029541Z",
+     "start_time": "2021-04-27T00:32:17.023938Z"
+    }
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "[3, 3, 3]"
+      ]
+     },
+     "execution_count": 29,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "list(x.shape)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 26,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2021-04-27T00:31:57.709249Z",
+     "start_time": "2021-04-27T00:31:57.701963Z"
+    }
+   },
+   "outputs": [
+    {
+     "ename": "AttributeError",
+     "evalue": "attribute 'shape' of 'torch._C._TensorBase' objects is not writable",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mAttributeError\u001b[0m                            Traceback (most recent call last)",
+      "\u001b[0;32m<ipython-input-26-2924dc973659>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mx\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mshape\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;36m5\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
+      "\u001b[0;31mAttributeError\u001b[0m: attribute 'shape' of 'torch._C._TensorBase' objects is not writable"
+     ]
+    }
+   ],
+   "source": [
+    "x.shape = 5"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 31,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2021-04-27T00:39:07.586468Z",
+     "start_time": "2021-04-27T00:39:07.581351Z"
+    }
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "[5]"
+      ]
+     },
+     "execution_count": 31,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "[1,2,3,4,5][4:]"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,