change initialize control of linear realNVP from True/False to scale

Author: Zhangyanbo

INN/CouplingModels/RealNVP/linear.py

@@ -4,14 +4,14 @@
 
 
 class NonlinearRealNVP(INNModule):
-    def __init__(self, dim=None, f_log_s=None, f_t=None, k=4, mask=None, clip=1, activation_fn=None):
+    def __init__(self, dim=None, f_log_s=None, f_t=None, k=4, mask=None, clip=1, activation_fn=None, scale=0.01):
         super(NonlinearRealNVP, self).__init__()
         self.dim = dim
 
         if f_log_s is None:
-            f_log_s = coupling_utils.default_nonlinear_net(dim, k, activation_fn, zero=True)
+            f_log_s = coupling_utils.default_nonlinear_net(dim, k, activation_fn, scale=scale)
         if f_t is None:
-            f_t = coupling_utils.default_nonlinear_net(dim, k, activation_fn)
+            f_t = coupling_utils.default_nonlinear_net(dim, k, activation_fn, scale=scale*10)
 
         self.net = utils.combined_real_nvp(dim, f_log_s, f_t, mask, clip)
 

INN/CouplingModels/utils.py

@@ -3,9 +3,9 @@
 
 
 class default_nonlinear_net(nn.Module):
-    def __init__(self, dim, k, activation_fn=None, zero=False):
+    def __init__(self, dim, k, activation_fn=None, scale=0.01):
         super(default_nonlinear_net, self).__init__()
-        self._zero_init = zero
+        self.scale = scale
         self.activation_fn = activation_fn
         self.net = self.default_net(dim, k, activation_fn)
 
@@ -25,10 +25,8 @@ def init_weights(self, m):
         if type(m) == nn.Linear:
             # doing xavier initialization
             # NOTE: Kaiming initialization will make the output too high, which leads to nan
-            if self._zero_init:
-                torch.nn.init.zeros_(m.weight.data)
-            else:
-                torch.nn.init.xavier_normal_(m.weight.data)
+            torch.nn.init.xavier_normal_(m.weight.data)
+            m.weight.data *= self.scale
             torch.nn.init.zeros_(m.bias.data)
 
     def forward(self, x):