code refactoring,

lr_dropout is now disabled in pretraining mode
changed help message for lr_dropout and random_warp

core/leras/archis.py

@@ -0,0 +1,565 @@
+
+def initialize_archis(nn):
+    tf = nn.tf
+    
+    def get_ae_models():
+        conv_kernel_initializer = nn.initializers.ca()
+
+        class Downscale(nn.ModelBase):
+            def __init__(self, in_ch, out_ch, kernel_size=5, dilations=1, subpixel=True, use_activator=True, *kwargs ):
+                self.in_ch = in_ch
+                self.out_ch = out_ch
+                self.kernel_size = kernel_size
+                self.dilations = dilations
+                self.subpixel = subpixel
+                self.use_activator = use_activator
+                super().__init__(*kwargs)
+
+            def on_build(self, *args, **kwargs ):
+                self.conv1 = nn.Conv2D( self.in_ch,
+                                        self.out_ch // (4 if self.subpixel else 1),
+                                        kernel_size=self.kernel_size,
+                                        strides=1 if self.subpixel else 2,
+                                        padding='SAME', dilations=self.dilations, kernel_initializer=conv_kernel_initializer)
+
+            def forward(self, x):
+                x = self.conv1(x)
+                if self.subpixel:
+                    x = nn.tf_space_to_depth(x, 2)
+                if self.use_activator:
+                    x = tf.nn.leaky_relu(x, 0.1)
+                return x
+
+            def get_out_ch(self):
+                return (self.out_ch // 4) * 4
+
+        class DownscaleBlock(nn.ModelBase):
+            def on_build(self, in_ch, ch, n_downscales, kernel_size, dilations=1, subpixel=True):
+                self.downs = []
+
+                last_ch = in_ch
+                for i in range(n_downscales):
+                    cur_ch = ch*( min(2**i, 8)  )
+                    self.downs.append ( Downscale(last_ch, cur_ch, kernel_size=kernel_size, dilations=dilations, subpixel=subpixel) )
+                    last_ch = self.downs[-1].get_out_ch()
+
+            def forward(self, inp):
+                x = inp
+                for down in self.downs:
+                    x = down(x)
+                return x
+
+        class Upscale(nn.ModelBase):
+            def on_build(self, in_ch, out_ch, kernel_size=3 ):
+                self.conv1 = nn.Conv2D( in_ch, out_ch*4, kernel_size=kernel_size, padding='SAME', kernel_initializer=conv_kernel_initializer)
+
+            def forward(self, x):
+                x = self.conv1(x)
+                x = tf.nn.leaky_relu(x, 0.1)
+                x = nn.tf_depth_to_space(x, 2)
+                return x
+
+        class ResidualBlock(nn.ModelBase):
+            def on_build(self, ch, kernel_size=3 ):
+                self.conv1 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME', kernel_initializer=conv_kernel_initializer)
+                self.conv2 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME', kernel_initializer=conv_kernel_initializer)
+
+            def forward(self, inp):
+                x = self.conv1(inp)
+                x = tf.nn.leaky_relu(x, 0.2)
+                x = self.conv2(x)
+                x = tf.nn.leaky_relu(inp + x, 0.2)
+                return x
+
+        class UpdownResidualBlock(nn.ModelBase):
+            def on_build(self, ch, inner_ch, kernel_size=3 ):
+                self.up   = Upscale (ch, inner_ch, kernel_size=kernel_size)
+                self.res  = ResidualBlock (inner_ch, kernel_size=kernel_size)
+                self.down = Downscale (inner_ch, ch, kernel_size=kernel_size, use_activator=False)
+
+            def forward(self, inp):
+                x = self.up(inp)
+                x = upx = self.res(x)
+                x = self.down(x)
+                x = x + inp
+                x = tf.nn.leaky_relu(x, 0.2)
+                return x, upx
+
+        class Encoder(nn.ModelBase):
+            def on_build(self, in_ch, e_ch, is_hd):
+                self.is_hd=is_hd
+                if self.is_hd:
+                    self.down1 = DownscaleBlock(in_ch, e_ch*2, n_downscales=4, kernel_size=3, dilations=1)
+                    self.down2 = DownscaleBlock(in_ch, e_ch*2, n_downscales=4, kernel_size=5, dilations=1)
+                    self.down3 = DownscaleBlock(in_ch, e_ch//2, n_downscales=4, kernel_size=5, dilations=2)
+                    self.down4 = DownscaleBlock(in_ch, e_ch//2, n_downscales=4, kernel_size=7, dilations=2)
+                else:
+                    self.down1 = DownscaleBlock(in_ch, e_ch, n_downscales=4, kernel_size=5, dilations=1, subpixel=False)
+
+            def forward(self, inp):
+                if self.is_hd:
+                    x = tf.concat([ nn.tf_flatten(self.down1(inp)),
+                                    nn.tf_flatten(self.down2(inp)),
+                                    nn.tf_flatten(self.down3(inp)),
+                                    nn.tf_flatten(self.down4(inp)) ], -1 )
+                else:
+                    x = nn.tf_flatten(self.down1(inp))
+                return x
+
+        class Inter(nn.ModelBase):
+            def __init__(self, in_ch, lowest_dense_res, ae_ch, ae_out_ch, is_hd=False, **kwargs):
+                self.in_ch, self.lowest_dense_res, self.ae_ch, self.ae_out_ch = in_ch, lowest_dense_res, ae_ch, ae_out_ch
+                super().__init__(**kwargs)
+
+            def on_build(self):
+                in_ch, lowest_dense_res, ae_ch, ae_out_ch = self.in_ch, self.lowest_dense_res, self.ae_ch, self.ae_out_ch
+
+                self.dense1 = nn.Dense( in_ch, ae_ch )
+                self.dense2 = nn.Dense( ae_ch, lowest_dense_res * lowest_dense_res * ae_out_ch )
+                self.upscale1 = Upscale(ae_out_ch, ae_out_ch)
+
+            def forward(self, inp):
+                x = self.dense1(inp)
+                x = self.dense2(x)
+                x = nn.tf_reshape_4D (x, self.lowest_dense_res, self.lowest_dense_res, self.ae_out_ch)
+                x = self.upscale1(x)
+                return x
+
+            def get_out_ch(self):
+                return self.ae_out_ch
+
+        class Decoder(nn.ModelBase):
+            def on_build(self, in_ch, d_ch, d_mask_ch, is_hd ):
+                self.is_hd = is_hd
+
+                self.upscale0 = Upscale(in_ch, d_ch*8, kernel_size=3)
+                self.upscale1 = Upscale(d_ch*8, d_ch*4, kernel_size=3)
+                self.upscale2 = Upscale(d_ch*4, d_ch*2, kernel_size=3)
+
+                if is_hd:
+                    self.res0 = UpdownResidualBlock(in_ch, d_ch*8, kernel_size=3)
+                    self.res1 = UpdownResidualBlock(d_ch*8, d_ch*4, kernel_size=3)
+                    self.res2 = UpdownResidualBlock(d_ch*4, d_ch*2, kernel_size=3)
+                    self.res3 = UpdownResidualBlock(d_ch*2, d_ch, kernel_size=3)
+                else:
+                    self.res0 = ResidualBlock(d_ch*8, kernel_size=3)
+                    self.res1 = ResidualBlock(d_ch*4, kernel_size=3)
+                    self.res2 = ResidualBlock(d_ch*2, kernel_size=3)
+
+                self.out_conv  = nn.Conv2D( d_ch*2, 3, kernel_size=1, padding='SAME', kernel_initializer=conv_kernel_initializer)
+
+                self.upscalem0 = Upscale(in_ch, d_mask_ch*8, kernel_size=3)
+                self.upscalem1 = Upscale(d_mask_ch*8, d_mask_ch*4, kernel_size=3)
+                self.upscalem2 = Upscale(d_mask_ch*4, d_mask_ch*2, kernel_size=3)
+                self.out_convm = nn.Conv2D( d_mask_ch*2, 1, kernel_size=1, padding='SAME', kernel_initializer=conv_kernel_initializer)
+
+            def get_weights_ex(self, include_mask):
+                # Call internal get_weights in order to initialize inner logic
+                self.get_weights()
+
+                weights = self.upscale0.get_weights() + self.upscale1.get_weights() + self.upscale2.get_weights() \
+                          + self.res0.get_weights() + self.res1.get_weights() + self.res2.get_weights() + self.out_conv.get_weights()
+                          
+                if self.is_hd:
+                    weights += self.res3.get_weights()
+
+                if include_mask:
+                    weights += self.upscalem0.get_weights() + self.upscalem1.get_weights() + self.upscalem2.get_weights() \
+                               + self.out_convm.get_weights()
+                return weights
+
+
+            def forward(self, inp):
+                z = inp
+
+                if self.is_hd:
+                    x, upx = self.res0(z)
+                    x = self.upscale0(x)
+                    x = tf.nn.leaky_relu(x + upx, 0.2)
+                    x, upx = self.res1(x)
+
+                    x = self.upscale1(x)
+                    x = tf.nn.leaky_relu(x + upx, 0.2)
+                    x, upx = self.res2(x)
+
+                    x = self.upscale2(x)
+                    x = tf.nn.leaky_relu(x + upx, 0.2)
+                    x, upx = self.res3(x)
+                else:
+                    x = self.upscale0(z)
+                    x = self.res0(x)
+                    x = self.upscale1(x)
+                    x = self.res1(x)
+                    x = self.upscale2(x)
+                    x = self.res2(x)
+
+                m = self.upscalem0(z)
+                m = self.upscalem1(m)
+                m = self.upscalem2(m)
+
+                return tf.nn.sigmoid(self.out_conv(x)), \
+                       tf.nn.sigmoid(self.out_convm(m))
+                       
+        
+                
+        return Encoder, Inter, Decoder
+    
+    nn.get_ae_models = get_ae_models
+        
+    def get_ae_models2():
+        conv_kernel_initializer = nn.initializers.ca()
+        class Downscale(nn.ModelBase):
+            def __init__(self, in_ch, out_ch, kernel_size=5, dilations=1, is_hd=False, *kwargs ):
+                self.in_ch = in_ch
+                self.out_ch = out_ch
+                self.kernel_size = kernel_size
+                self.dilations = dilations
+                self.is_hd = is_hd
+                super().__init__(*kwargs)
+
+            def on_build(self, *args, **kwargs ):
+                if not self.is_hd:
+                    self.conv1 = nn.Conv2D( self.in_ch, self.out_ch, kernel_size=self.kernel_size, strides=2, padding='SAME', dilations=self.dilations, kernel_initializer=conv_kernel_initializer)
+                else:
+                    self.conv0 = nn.Conv2D( self.in_ch, self.out_ch//4, kernel_size=self.kernel_size, strides=2, padding='SAME', dilations=self.dilations, kernel_initializer=conv_kernel_initializer)
+                    self.conv1 = nn.Conv2D( self.out_ch//4, self.out_ch//4, kernel_size=self.kernel_size, strides=1, padding='SAME', dilations=self.dilations, kernel_initializer=conv_kernel_initializer)
+                    self.conv2 = nn.Conv2D( self.out_ch//4, self.out_ch//4, kernel_size=self.kernel_size, strides=1, padding='SAME', dilations=self.dilations, kernel_initializer=conv_kernel_initializer)
+                    self.conv3 = nn.Conv2D( self.out_ch//4, self.out_ch//4, kernel_size=self.kernel_size, strides=1, padding='SAME', dilations=self.dilations, kernel_initializer=conv_kernel_initializer)
+                
+
+            def forward(self, x):
+                if not self.is_hd:
+                    x = self.conv1(x)
+                else:
+                    x = x0 = self.conv0(x)
+                    x = x1 = self.conv1(x)
+                    x = x2 = self.conv2(x)
+                    x = x3 = self.conv3(x)
+                    x = tf.concat([x0,x1,x2,x3], axis=nn.conv2d_ch_axis)                    
+                x = tf.nn.leaky_relu(x, 0.1)                
+                return x
+
+            def get_out_ch(self):
+                return (self.out_ch // 4) * 4
+
+        class DownscaleBlock(nn.ModelBase):
+            def on_build(self, in_ch, ch, n_downscales, kernel_size, dilations=1, is_hd=False):
+                self.downs = []
+
+                last_ch = in_ch
+                for i in range(n_downscales):
+                    cur_ch = ch*( min(2**i, 8)  )
+                    self.downs.append ( Downscale(last_ch, cur_ch, kernel_size=kernel_size, dilations=dilations, is_hd=is_hd) )
+                    last_ch = self.downs[-1].get_out_ch()
+
+            def forward(self, inp):
+                x = inp
+                for down in self.downs:
+                    x = down(x)
+                return x
+        
+        class Upscale(nn.ModelBase):
+            def on_build(self, in_ch, out_ch, kernel_size=3, is_hd=False ):
+                self.is_hd = is_hd
+                
+                if not self.is_hd:
+                    self.conv1 = nn.Conv2D( in_ch, out_ch*4, kernel_size=kernel_size, padding='SAME', kernel_initializer=conv_kernel_initializer)
+                else:
+                    self.conv0 = nn.Conv2D( in_ch, out_ch, kernel_size=kernel_size, padding='SAME', kernel_initializer=conv_kernel_initializer)
+                    self.conv1 = nn.Conv2D( out_ch, out_ch, kernel_size=kernel_size, padding='SAME', kernel_initializer=conv_kernel_initializer)
+                    self.conv2 = nn.Conv2D( out_ch, out_ch, kernel_size=kernel_size, padding='SAME', kernel_initializer=conv_kernel_initializer)
+                    self.conv3 = nn.Conv2D( out_ch, out_ch, kernel_size=kernel_size, padding='SAME', kernel_initializer=conv_kernel_initializer)
+                    
+            def forward(self, x):
+                if not self.is_hd:
+                    x = self.conv1(x)
+                else:
+                    x = x0 = self.conv0(x)
+                    x = x1 = self.conv1(x)
+                    x = x2 = self.conv2(x)
+                    x = x3 = self.conv3(x)                    
+                    x = tf.concat([x0,x1,x2,x3], axis=nn.conv2d_ch_axis)
+                    
+                x = tf.nn.leaky_relu(x, 0.1)
+                x = nn.tf_depth_to_space(x, 2)
+                return x
+                        
+        class ResidualBlock(nn.ModelBase):
+            def on_build(self, ch, kernel_size=3, is_hd=False ):
+                self.is_hd = is_hd
+                if not is_hd:
+                    self.conv1 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME', kernel_initializer=conv_kernel_initializer)
+                    self.conv2 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME', kernel_initializer=conv_kernel_initializer)
+                else:
+                    self.conv10 = nn.Conv2D( ch, ch//4, kernel_size=kernel_size, padding='SAME', kernel_initializer=conv_kernel_initializer)
+                    self.conv11 = nn.Conv2D( ch//4, ch//4, kernel_size=kernel_size, padding='SAME', kernel_initializer=conv_kernel_initializer)
+                    self.conv12 = nn.Conv2D( ch//4, ch//4, kernel_size=kernel_size, padding='SAME', kernel_initializer=conv_kernel_initializer)
+                    self.conv13 = nn.Conv2D( ch//4, ch//4, kernel_size=kernel_size, padding='SAME', kernel_initializer=conv_kernel_initializer)
+                    
+                    self.conv2 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME', kernel_initializer=conv_kernel_initializer)
+
+            def forward(self, inp):
+                if not self.is_hd:
+                    x = self.conv1(inp)
+                else:
+                    x = x0 = self.conv10(inp)
+                    x = x1 = self.conv11(x)
+                    x = x2 = self.conv12(x)
+                    x = x3 = self.conv13(x)                    
+                    x = tf.concat([x0,x1,x2,x3], axis=nn.conv2d_ch_axis)                    
+                    
+                x = tf.nn.leaky_relu(x, 0.2)
+                x = self.conv2(x)
+                x = tf.nn.leaky_relu(inp + x, 0.2)
+                return x
+
+        class Encoder(nn.ModelBase):
+            def on_build(self, in_ch, e_ch, is_hd):
+                self.is_hd=is_hd
+                if self.is_hd:
+                    self.down1 = DownscaleBlock(in_ch, e_ch*2, n_downscales=4, kernel_size=3, is_hd=is_hd)
+                    self.down2 = DownscaleBlock(in_ch, e_ch*2, n_downscales=4, kernel_size=5, is_hd=is_hd)
+                else:
+                    self.down1 = DownscaleBlock(in_ch, e_ch, n_downscales=4, kernel_size=5, is_hd=is_hd)
+
+            def forward(self, inp):
+                if self.is_hd:
+                    x = tf.concat([ nn.tf_flatten(self.down1(inp)),
+                                    nn.tf_flatten(self.down2(inp)) ], -1 )
+                else:
+                    x = nn.tf_flatten(self.down1(inp))
+                return x
+
+        class Inter(nn.ModelBase):
+            def __init__(self, in_ch, lowest_dense_res, ae_ch, ae_out_ch, is_hd=False, **kwargs):
+                self.in_ch, self.lowest_dense_res, self.ae_ch, self.ae_out_ch, self.is_hd = in_ch, lowest_dense_res, ae_ch, ae_out_ch, is_hd
+                super().__init__(**kwargs)
+
+            def on_build(self):
+                in_ch, lowest_dense_res, ae_ch, ae_out_ch = self.in_ch, self.lowest_dense_res, self.ae_ch, self.ae_out_ch
+
+                self.dense1 = nn.Dense( in_ch, ae_ch )
+                self.dense2 = nn.Dense( ae_ch, lowest_dense_res * lowest_dense_res * ae_out_ch )
+                self.upscale1 = Upscale(ae_out_ch, ae_out_ch, is_hd=self.is_hd)
+
+            def forward(self, inp):
+                x = self.dense1(inp)
+                x = self.dense2(x)
+                x = nn.tf_reshape_4D (x, self.lowest_dense_res, self.lowest_dense_res, self.ae_out_ch)
+                x = self.upscale1(x)
+                return x
+
+            def get_out_ch(self):
+                return self.ae_out_ch
+
+        class Decoder(nn.ModelBase):
+            def on_build(self, in_ch, d_ch, d_mask_ch, is_hd ):
+                self.is_hd = is_hd
+
+                self.upscale0 = Upscale(in_ch, d_ch*8, kernel_size=3)
+                self.upscale1 = Upscale(d_ch*8, d_ch*4, kernel_size=3)
+                self.upscale2 = Upscale(d_ch*4, d_ch*2, kernel_size=3)
+
+                self.res0 = ResidualBlock(d_ch*8, kernel_size=3, is_hd=is_hd)
+                self.res1 = ResidualBlock(d_ch*4, kernel_size=3, is_hd=is_hd)
+                self.res2 = ResidualBlock(d_ch*2, kernel_size=3, is_hd=is_hd)
+
+                self.out_conv  = nn.Conv2D( d_ch*2, 3, kernel_size=1, padding='SAME', kernel_initializer=conv_kernel_initializer)
+
+                self.upscalem0 = Upscale(in_ch, d_mask_ch*8, kernel_size=3)
+                self.upscalem1 = Upscale(d_mask_ch*8, d_mask_ch*4, kernel_size=3)
+                self.upscalem2 = Upscale(d_mask_ch*4, d_mask_ch*2, kernel_size=3)
+                self.out_convm = nn.Conv2D( d_mask_ch*2, 1, kernel_size=1, padding='SAME', kernel_initializer=conv_kernel_initializer)
+
+            def get_weights_ex(self, include_mask):
+                # Call internal get_weights in order to initialize inner logic
+                self.get_weights()
+
+                weights = self.upscale0.get_weights() + self.upscale1.get_weights() + self.upscale2.get_weights() \
+                          + self.res0.get_weights() + self.res1.get_weights() + self.res2.get_weights() + self.out_conv.get_weights()
+
+                if include_mask:
+                    weights += self.upscalem0.get_weights() + self.upscalem1.get_weights() + self.upscalem2.get_weights() \
+                               + self.out_convm.get_weights()
+                return weights
+
+
+            def forward(self, inp):
+                z = inp
+                
+                x = self.upscale0(z)
+                x = self.res0(x)
+                x = self.upscale1(x)
+                x = self.res1(x)
+                x = self.upscale2(x)
+                x = self.res2(x)
+
+                m = self.upscalem0(z)
+                m = self.upscalem1(m)
+                m = self.upscalem2(m)
+
+                return tf.nn.sigmoid(self.out_conv(x)), \
+                       tf.nn.sigmoid(self.out_convm(m))
+                       
+                
+        return Encoder, Inter, Decoder
+                
+    nn.get_ae_models2 = get_ae_models2
+    
+    def get_ae_models_chervoniy():
+        """
+        by @chervoniy
+        """
+        conv_kernel_initializer = nn.initializers.ca()
+
+        class Downscale(nn.ModelBase):
+            def __init__(self, in_ch, kernel_size=3, dilations=1, *kwargs ):
+                self.in_ch = in_ch
+                self.kernel_size = kernel_size
+                self.dilations = dilations
+                super().__init__(*kwargs)
+
+            def on_build(self, *args, **kwargs ):
+                self.conv_base1 = nn.Conv2D( self.in_ch, self.in_ch//2, kernel_size=1, strides=1, padding='SAME', dilations=self.dilations, kernel_initializer=conv_kernel_initializer)
+                self.conv_l1 = nn.Conv2D( self.in_ch//2, self.in_ch//2, kernel_size=self.kernel_size, strides=1, padding='SAME', dilations=self.dilations, kernel_initializer=conv_kernel_initializer)
+                self.conv_l2 = nn.Conv2D( self.in_ch//2, self.in_ch//2, kernel_size=self.kernel_size, strides=2, padding='SAME', dilations=self.dilations, kernel_initializer=conv_kernel_initializer)
+
+                self.conv_base2 = nn.Conv2D( self.in_ch, self.in_ch//2, kernel_size=1, strides=1, padding='SAME', dilations=self.dilations, kernel_initializer=conv_kernel_initializer)
+                self.conv_r1 = nn.Conv2D( self.in_ch//2, self.in_ch//2, kernel_size=self.kernel_size, strides=2, padding='SAME', dilations=self.dilations, kernel_initializer=conv_kernel_initializer)
+
+                self.pool_size = [1,1,2,2] if nn.data_format == 'NCHW' else [1,2,2,1]
+            def forward(self, x):
+
+                x_l = self.conv_base1(x)
+                x_l = self.conv_l1(x_l)
+                x_l = self.conv_l2(x_l)
+
+                x_r = self.conv_base2(x)
+                x_r = self.conv_r1(x_r)
+
+                x_pool = tf.nn.max_pool(x, ksize=self.pool_size, strides=self.pool_size, padding='SAME', data_format=nn.data_format)
+
+                x = tf.concat([x_l, x_r, x_pool], axis=nn.conv2d_ch_axis)
+                x = nn.tf_gelu(x)
+                return x
+
+        class Upscale(nn.ModelBase):
+            def on_build(self, in_ch, out_ch, kernel_size=3 ):
+                self.conv1 = nn.Conv2D( in_ch, out_ch, kernel_size=kernel_size, padding='SAME', kernel_initializer=conv_kernel_initializer)
+                self.conv2 = nn.Conv2D( out_ch, out_ch, kernel_size=kernel_size, padding='SAME', kernel_initializer=conv_kernel_initializer)
+                self.conv3 = nn.Conv2D( out_ch, out_ch, kernel_size=kernel_size, padding='SAME', kernel_initializer=conv_kernel_initializer)
+                self.conv4 = nn.Conv2D( out_ch, out_ch, kernel_size=kernel_size, padding='SAME', kernel_initializer=conv_kernel_initializer)
+
+            def forward(self, x):
+                x0 = self.conv1(x)
+                x1 = self.conv2(x0)
+                x2 = self.conv3(x1)
+                x3 = self.conv4(x2)
+                x = tf.concat([x0, x1, x2, x3], axis=nn.conv2d_ch_axis)
+                x = nn.tf_gelu(x)
+                x = nn.tf_depth_to_space(x, 2)
+                return x
+
+        class ResidualBlock(nn.ModelBase):
+            def on_build(self, ch, kernel_size=3 ):
+                self.conv1 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME', kernel_initializer=conv_kernel_initializer)
+                self.conv2 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME', kernel_initializer=conv_kernel_initializer)
+
+            def forward(self, inp):
+                x = self.conv1(inp)
+                x = tf.nn.leaky_relu(x, 0.2)
+                x = self.conv2(x)
+                x = tf.nn.leaky_relu(inp + x, 0.2)
+                return x
+
+        class Encoder(nn.ModelBase):
+            def on_build(self, in_ch, e_ch):
+                self.conv0 = nn.Conv2D(in_ch, e_ch, kernel_size=3, padding='SAME', kernel_initializer=conv_kernel_initializer)
+
+                self.down0 = Downscale(e_ch)
+                self.down1 = Downscale(e_ch*2)
+                self.down2 = Downscale(e_ch*4)
+                self.down3 = Downscale(e_ch*8)
+                self.down4 = Downscale(e_ch*16)
+
+            def forward(self, inp):
+                x = self.conv0(inp)
+                x = self.down0(x)
+                x = self.down1(x)
+                x = self.down2(x)
+                x = self.down3(x)
+                x = self.down4(x)
+                x = nn.tf_flatten(x)
+                return x
+
+        class Inter(nn.ModelBase):
+            def __init__(self, in_ch, lowest_dense_res, ae_ch, ae_out_ch, **kwargs):
+                self.in_ch, self.lowest_dense_res, self.ae_ch, self.ae_out_ch = in_ch, lowest_dense_res, ae_ch, ae_out_ch
+                super().__init__(**kwargs)
+
+            def on_build(self):
+                in_ch, lowest_dense_res, ae_ch, ae_out_ch = self.in_ch, self.lowest_dense_res, self.ae_ch, self.ae_out_ch
+
+                self.dense_l = nn.Dense( in_ch, ae_ch//2, kernel_initializer=tf.initializers.orthogonal)
+                self.dense_r = nn.Dense( in_ch, ae_ch//2, maxout_features=4, kernel_initializer=tf.initializers.orthogonal)
+                self.dense = nn.Dense( ae_ch, lowest_dense_res * lowest_dense_res * (ae_out_ch//2), kernel_initializer=tf.initializers.orthogonal)
+                self.upscale1 = Upscale(ae_out_ch//2, ae_out_ch//2)
+
+            def forward(self, inp):
+                x0 = self.dense_l(inp)
+                x1 = self.dense_r(inp)
+                x = tf.concat([x0, x1], axis=-1)
+                x = self.dense(x)
+                x = nn.tf_reshape_4D (x, lowest_dense_res, lowest_dense_res, self.ae_out_ch//2)
+                x = self.upscale1(x)
+                return x
+
+            def get_out_ch(self):
+                return self.ae_out_ch//2
+
+        class Decoder(nn.ModelBase):
+            def on_build(self, in_ch, d_ch, d_mask_ch ):
+
+                self.upscale0 = Upscale(in_ch, d_ch*8)
+                self.upscale1 = Upscale(d_ch*8, d_ch*4)
+                self.upscale2 = Upscale(d_ch*4, d_ch*2)
+                self.upscale3 = Upscale(d_ch*2, d_ch)
+
+                self.out_conv  = nn.Conv2D( d_ch, 3, kernel_size=1, padding='SAME', kernel_initializer=conv_kernel_initializer)
+
+                self.upscalem0 = Upscale(in_ch, d_mask_ch*8, kernel_size=3)
+                self.upscalem1 = Upscale(d_mask_ch*8, d_mask_ch*4, kernel_size=3)
+                self.upscalem2 = Upscale(d_mask_ch*4, d_mask_ch*2, kernel_size=3)
+                self.upscalem3 = Upscale(d_mask_ch*2, d_mask_ch, kernel_size=3)
+                self.out_convm = nn.Conv2D( d_mask_ch, 1, kernel_size=1, padding='SAME', kernel_initializer=conv_kernel_initializer)
+
+            def get_weights_ex(self, include_mask):
+                # Call internal get_weights in order to initialize inner logic
+                self.get_weights()
+
+                weights = self.upscale0.get_weights() + self.upscale1.get_weights() + self.upscale2.get_weights() + self.upscale3.get_weights() + self.out_conv.get_weights()
+
+                if include_mask:
+                    weights += self.upscalem0.get_weights() + self.upscalem1.get_weights() + self.upscalem2.get_weights() + self.upscale3.get_weights() + self.out_convm.get_weights()
+                return weights
+
+
+            def forward(self, inp):
+                z = inp
+
+                x = self.upscale0(inp)
+                x = self.upscale1(x)
+                x = self.upscale2(x)
+                x = self.upscale3(x)
+
+                m = self.upscalem0(z)
+                m = self.upscalem1(m)
+                m = self.upscalem2(m)
+                m = self.upscalem3(m)
+
+                return tf.nn.sigmoid(self.out_conv(x)), \
+                       tf.nn.sigmoid(self.out_convm(m))
+
+        return Encoder, Inter, Decoder
+        
+    return get_ae_models_chervoniy

core/leras/models.py

@@ -1,4 +1,6 @@
 import types
+import numpy as np
+from core.interact import interact as io
 
 def initialize_models(nn):
     tf = nn.tf
@@ -296,6 +298,28 @@ def initialize_models(nn):
 
     nn.IllumDiscriminator = IllumDiscriminator
     
+    class CodeDiscriminator(nn.ModelBase):
+        def on_build(self, in_ch, code_res, ch=256, conv_kernel_initializer=None):
+            if conv_kernel_initializer is None:
+                conv_kernel_initializer = nn.initializers.ca()
+                
+            n_downscales = 1 + code_res // 8
+
+            self.convs = []
+            prev_ch = in_ch
+            for i in range(n_downscales):
+                cur_ch = ch * min( (2**i), 8 )
+                self.convs.append ( nn.Conv2D( prev_ch, cur_ch, kernel_size=4 if i == 0 else 3, strides=2, padding='SAME', kernel_initializer=conv_kernel_initializer) )
+                prev_ch = cur_ch
+
+            self.out_conv =  nn.Conv2D( prev_ch, 1, kernel_size=1, padding='VALID', kernel_initializer=conv_kernel_initializer)
+
+        def forward(self, x):
+            for conv in self.convs:
+                x = tf.nn.leaky_relu( conv(x), 0.1 )
+            return self.out_conv(x)
+    nn.CodeDiscriminator = CodeDiscriminator
+            
 patch_discriminator_kernels = \
     { 1  : (512, [ [1,1] ]),
       2  : (512, [ [2,1] ]),

core/leras/nn.py

@@ -86,6 +86,7 @@ class nn():
     # Models
     PatchDiscriminator = None
     IllumDiscriminator = None
+    CodeDiscriminator = None
     
     @staticmethod
     def initialize(device_config=None, floatx="float32", data_format="NHWC"):
@@ -147,12 +148,14 @@ class nn():
             from .initializers import initialize_initializers
             from .optimizers import initialize_optimizers
             from .models import initialize_models
+            from .archis import initialize_archis
             
             initialize_tensor_ops(nn)
             initialize_layers(nn)
             initialize_initializers(nn)
             initialize_optimizers(nn)
             initialize_models(nn)
+            initialize_archis(nn)
             
         if nn.tf_sess is None:
             nn.tf_sess = tf.Session(config=nn.tf_sess_config)

models/Model_SAEHD/Model.py

@@ -94,8 +94,8 @@ class SAEHDModel(ModelBase):
         if self.is_first_run() or ask_override:
             self.options['models_opt_on_gpu'] = io.input_bool ("Place models and optimizer on GPU", default_models_opt_on_gpu, help_message="When you train on one GPU, by default model and optimizer weights are placed on GPU to accelerate the process. You can place they on CPU to free up extra VRAM, thus set bigger dimensions.")
 
-            self.options['lr_dropout']  = io.input_bool ("Use learning rate dropout", default_lr_dropout, help_message="When the face is trained enough, you can enable this option to get extra sharpness for less amount of iterations.")
+            self.options['lr_dropout']  = io.input_bool ("Use learning rate dropout", default_lr_dropout, help_message="When the face is trained enough, you can enable this option to get extra sharpness and reduce subpixel shake for less amount of iterations.")
-            self.options['random_warp'] = io.input_bool ("Enable random warp of samples", default_random_warp, help_message="Random warp is required to generalize facial expressions of both faces. When the face is trained enough, you can disable it to get extra sharpness for less amount of iterations.")
+            self.options['random_warp'] = io.input_bool ("Enable random warp of samples", default_random_warp, help_message="Random warp is required to generalize facial expressions of both faces. When the face is trained enough, you can disable it to get extra sharpness and reduce subpixel shake for less amount of iterations.")
 
             self.options['gan_power'] = np.clip ( io.input_number ("GAN power", default_gan_power, add_info="0.0 .. 10.0", help_message="Train the network in Generative Adversarial manner. Accelerates the speed of training. Forces the neural network to learn small details of the face. You can enable/disable this option at any time. Typical value is 1.0"), 0.0, 10.0 )
 
@@ -125,220 +125,7 @@ class SAEHDModel(ModelBase):
         nn.initialize(data_format=self.model_data_format)
         tf = nn.tf
 
-        conv_kernel_initializer = nn.initializers.ca()
+        Encoder, Inter, Decoder = nn.get_ae_models()
-
-        class Downscale(nn.ModelBase):
-            def __init__(self, in_ch, out_ch, kernel_size=5, dilations=1, subpixel=True, use_activator=True, *kwargs ):
-                self.in_ch = in_ch
-                self.out_ch = out_ch
-                self.kernel_size = kernel_size
-                self.dilations = dilations
-                self.subpixel = subpixel
-                self.use_activator = use_activator
-                super().__init__(*kwargs)
-
-            def on_build(self, *args, **kwargs ):
-                self.conv1 = nn.Conv2D( self.in_ch,
-                                        self.out_ch // (4 if self.subpixel else 1),
-                                        kernel_size=self.kernel_size,
-                                        strides=1 if self.subpixel else 2,
-                                        padding='SAME', dilations=self.dilations, kernel_initializer=conv_kernel_initializer)
-
-            def forward(self, x):
-                x = self.conv1(x)
-                if self.subpixel:
-                    x = nn.tf_space_to_depth(x, 2)
-                if self.use_activator:
-                    x = tf.nn.leaky_relu(x, 0.1)
-                return x
-
-            def get_out_ch(self):
-                return (self.out_ch // 4) * 4
-
-        class DownscaleBlock(nn.ModelBase):
-            def on_build(self, in_ch, ch, n_downscales, kernel_size, dilations=1, subpixel=True):
-                self.downs = []
-
-                last_ch = in_ch
-                for i in range(n_downscales):
-                    cur_ch = ch*( min(2**i, 8)  )
-                    self.downs.append ( Downscale(last_ch, cur_ch, kernel_size=kernel_size, dilations=dilations, subpixel=subpixel) )
-                    last_ch = self.downs[-1].get_out_ch()
-
-            def forward(self, inp):
-                x = inp
-                for down in self.downs:
-                    x = down(x)
-                return x
-
-        class Upscale(nn.ModelBase):
-            def on_build(self, in_ch, out_ch, kernel_size=3 ):
-                self.conv1 = nn.Conv2D( in_ch, out_ch*4, kernel_size=kernel_size, padding='SAME', kernel_initializer=conv_kernel_initializer)
-
-            def forward(self, x):
-                x = self.conv1(x)
-                x = tf.nn.leaky_relu(x, 0.1)
-                x = nn.tf_depth_to_space(x, 2)
-                return x
-
-        class ResidualBlock(nn.ModelBase):
-            def on_build(self, ch, kernel_size=3 ):
-                self.conv1 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME', kernel_initializer=conv_kernel_initializer)
-                self.conv2 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME', kernel_initializer=conv_kernel_initializer)
-
-            def forward(self, inp):
-                x = self.conv1(inp)
-                x = tf.nn.leaky_relu(x, 0.2)
-                x = self.conv2(x)
-                x = tf.nn.leaky_relu(inp + x, 0.2)
-                return x
-
-        class UpdownResidualBlock(nn.ModelBase):
-            def on_build(self, ch, inner_ch, kernel_size=3 ):
-                self.up   = Upscale (ch, inner_ch, kernel_size=kernel_size)
-                self.res  = ResidualBlock (inner_ch, kernel_size=kernel_size)
-                self.down = Downscale (inner_ch, ch, kernel_size=kernel_size, use_activator=False)
-
-            def forward(self, inp):
-                x = self.up(inp)
-                x = upx = self.res(x)
-                x = self.down(x)
-                x = x + inp
-                x = tf.nn.leaky_relu(x, 0.2)
-                return x, upx
-
-        class Encoder(nn.ModelBase):
-            def on_build(self, in_ch, e_ch, is_hd):
-                self.is_hd=is_hd
-                if self.is_hd:
-                    self.down1 = DownscaleBlock(in_ch, e_ch*2, n_downscales=4, kernel_size=3, dilations=1)
-                    self.down2 = DownscaleBlock(in_ch, e_ch*2, n_downscales=4, kernel_size=5, dilations=1)
-                    self.down3 = DownscaleBlock(in_ch, e_ch//2, n_downscales=4, kernel_size=5, dilations=2)
-                    self.down4 = DownscaleBlock(in_ch, e_ch//2, n_downscales=4, kernel_size=7, dilations=2)
-                else:
-                    self.down1 = DownscaleBlock(in_ch, e_ch, n_downscales=4, kernel_size=5, dilations=1, subpixel=False)
-
-            def forward(self, inp):
-                if self.is_hd:
-                    x = tf.concat([ nn.tf_flatten(self.down1(inp)),
-                                    nn.tf_flatten(self.down2(inp)),
-                                    nn.tf_flatten(self.down3(inp)),
-                                    nn.tf_flatten(self.down4(inp)) ], -1 )
-                else:
-                    x = nn.tf_flatten(self.down1(inp))
-                return x
-
-        class Inter(nn.ModelBase):
-            def __init__(self, in_ch, lowest_dense_res, ae_ch, ae_out_ch, **kwargs):
-                self.in_ch, self.lowest_dense_res, self.ae_ch, self.ae_out_ch = in_ch, lowest_dense_res, ae_ch, ae_out_ch
-                super().__init__(**kwargs)
-
-            def on_build(self):
-                in_ch, lowest_dense_res, ae_ch, ae_out_ch = self.in_ch, self.lowest_dense_res, self.ae_ch, self.ae_out_ch
-
-                self.dense1 = nn.Dense( in_ch, ae_ch )
-                self.dense2 = nn.Dense( ae_ch, lowest_dense_res * lowest_dense_res * ae_out_ch )
-                self.upscale1 = Upscale(ae_out_ch, ae_out_ch)
-
-            def forward(self, inp):
-                x = self.dense1(inp)
-                x = self.dense2(x)
-                x = nn.tf_reshape_4D (x, lowest_dense_res, lowest_dense_res, self.ae_out_ch)
-                x = self.upscale1(x)
-                return x
-
-            def get_out_ch(self):
-                return self.ae_out_ch
-
-        class Decoder(nn.ModelBase):
-            def on_build(self, in_ch, d_ch, d_mask_ch, is_hd ):
-                self.is_hd = is_hd
-
-                self.upscale0 = Upscale(in_ch, d_ch*8, kernel_size=3)
-                self.upscale1 = Upscale(d_ch*8, d_ch*4, kernel_size=3)
-                self.upscale2 = Upscale(d_ch*4, d_ch*2, kernel_size=3)
-
-                if is_hd:
-                    self.res0 = UpdownResidualBlock(in_ch, d_ch*8, kernel_size=3)
-                    self.res1 = UpdownResidualBlock(d_ch*8, d_ch*4, kernel_size=3)
-                    self.res2 = UpdownResidualBlock(d_ch*4, d_ch*2, kernel_size=3)
-                    self.res3 = UpdownResidualBlock(d_ch*2, d_ch, kernel_size=3)
-                else:
-                    self.res0 = ResidualBlock(d_ch*8, kernel_size=3)
-                    self.res1 = ResidualBlock(d_ch*4, kernel_size=3)
-                    self.res2 = ResidualBlock(d_ch*2, kernel_size=3)
-
-                self.out_conv  = nn.Conv2D( d_ch*2, 3, kernel_size=1, padding='SAME', kernel_initializer=conv_kernel_initializer)
-
-                self.upscalem0 = Upscale(in_ch, d_mask_ch*8, kernel_size=3)
-                self.upscalem1 = Upscale(d_mask_ch*8, d_mask_ch*4, kernel_size=3)
-                self.upscalem2 = Upscale(d_mask_ch*4, d_mask_ch*2, kernel_size=3)
-                self.out_convm = nn.Conv2D( d_mask_ch*2, 1, kernel_size=1, padding='SAME', kernel_initializer=conv_kernel_initializer)
-
-            def get_weights_ex(self, include_mask):
-                # Call internal get_weights in order to initialize inner logic
-                self.get_weights()
-
-                weights = self.upscale0.get_weights() + self.upscale1.get_weights() + self.upscale2.get_weights() \
-                          + self.res0.get_weights() + self.res1.get_weights() + self.res2.get_weights() + self.out_conv.get_weights()
-                          
-                if self.is_hd:
-                    weights += self.res3.get_weights()
-
-                if include_mask:
-                    weights += self.upscalem0.get_weights() + self.upscalem1.get_weights() + self.upscalem2.get_weights() \
-                               + self.out_convm.get_weights()
-                return weights
-
-
-            def forward(self, inp):
-                z = inp
-
-                if self.is_hd:
-                    x, upx = self.res0(z)
-                    x = self.upscale0(x)
-                    x = tf.nn.leaky_relu(x + upx, 0.2)
-                    x, upx = self.res1(x)
-
-                    x = self.upscale1(x)
-                    x = tf.nn.leaky_relu(x + upx, 0.2)
-                    x, upx = self.res2(x)
-
-                    x = self.upscale2(x)
-                    x = tf.nn.leaky_relu(x + upx, 0.2)
-                    x, upx = self.res3(x)
-                else:
-                    x = self.upscale0(z)
-                    x = self.res0(x)
-                    x = self.upscale1(x)
-                    x = self.res1(x)
-                    x = self.upscale2(x)
-                    x = self.res2(x)
-
-                m = self.upscalem0(z)
-                m = self.upscalem1(m)
-                m = self.upscalem2(m)
-
-                return tf.nn.sigmoid(self.out_conv(x)), \
-                       tf.nn.sigmoid(self.out_convm(m))
-
-        class CodeDiscriminator(nn.ModelBase):
-            def on_build(self, in_ch, code_res, ch=256):
-                n_downscales = 1 + code_res // 8
-
-                self.convs = []
-                prev_ch = in_ch
-                for i in range(n_downscales):
-                    cur_ch = ch * min( (2**i), 8 )
-                    self.convs.append ( nn.Conv2D( prev_ch, cur_ch, kernel_size=4 if i == 0 else 3, strides=2, padding='SAME', kernel_initializer=conv_kernel_initializer) )
-                    prev_ch = cur_ch
-
-                self.out_conv =  nn.Conv2D( prev_ch, 1, kernel_size=1, padding='VALID', kernel_initializer=conv_kernel_initializer)
-
-            def forward(self, x):
-                for conv in self.convs:
-                    x = tf.nn.leaky_relu( conv(x), 0.1 )
-                return self.out_conv(x)
   
         device_config = nn.getCurrentDeviceConfig()
         devices = device_config.devices
@@ -352,6 +139,7 @@ class SAEHDModel(ModelBase):
         learn_mask = self.options['learn_mask']
         eyes_prio = self.options['eyes_prio']
         archi = self.options['archi']
+        is_hd = 'hd' in archi
         ae_dims = self.options['ae_dims']
         e_dims = self.options['e_dims']
         d_dims = self.options['d_dims']
@@ -394,14 +182,14 @@ class SAEHDModel(ModelBase):
         # Initializing model classes
         with tf.device (models_opt_device):
             if 'df' in archi:
-                self.encoder = Encoder(in_ch=input_ch, e_ch=e_dims, is_hd='hd' in archi, name='encoder')
+                self.encoder = Encoder(in_ch=input_ch, e_ch=e_dims, is_hd=is_hd, name='encoder')
                 encoder_out_ch = self.encoder.compute_output_channels ( (nn.tf_floatx, bgr_shape))
 
-                self.inter = Inter (in_ch=encoder_out_ch, lowest_dense_res=lowest_dense_res, ae_ch=ae_dims, ae_out_ch=ae_dims, name='inter')
+                self.inter = Inter (in_ch=encoder_out_ch, lowest_dense_res=lowest_dense_res, ae_ch=ae_dims, ae_out_ch=ae_dims, is_hd=is_hd, name='inter')
                 inter_out_ch = self.inter.compute_output_channels ( (nn.tf_floatx, (None,encoder_out_ch)))
 
-                self.decoder_src = Decoder(in_ch=inter_out_ch, d_ch=d_dims, d_mask_ch=d_mask_dims, is_hd='hd' in archi, name='decoder_src')
+                self.decoder_src = Decoder(in_ch=inter_out_ch, d_ch=d_dims, d_mask_ch=d_mask_dims, is_hd=is_hd, name='decoder_src')
-                self.decoder_dst = Decoder(in_ch=inter_out_ch, d_ch=d_dims, d_mask_ch=d_mask_dims, is_hd='hd' in archi, name='decoder_dst')
+                self.decoder_dst = Decoder(in_ch=inter_out_ch, d_ch=d_dims, d_mask_ch=d_mask_dims, is_hd=is_hd, name='decoder_dst')
 
                 self.model_filename_list += [ [self.encoder,     'encoder.npy'    ],
                                               [self.inter,       'inter.npy'      ],
@@ -410,20 +198,20 @@ class SAEHDModel(ModelBase):
 
                 if self.is_training:
                     if self.options['true_face_power'] != 0:
-                        self.code_discriminator = CodeDiscriminator(ae_dims, code_res=lowest_dense_res*2, name='dis' )
+                        self.code_discriminator = nn.CodeDiscriminator(ae_dims, code_res=lowest_dense_res*2, name='dis' )
                         self.model_filename_list += [ [self.code_discriminator, 'code_discriminator.npy'] ]
 
             elif 'liae' in archi:
-                self.encoder = Encoder(in_ch=input_ch, e_ch=e_dims, is_hd='hd' in archi, name='encoder')
+                self.encoder = Encoder(in_ch=input_ch, e_ch=e_dims, is_hd=is_hd, name='encoder')
                 encoder_out_ch = self.encoder.compute_output_channels ( (nn.tf_floatx, bgr_shape))
 
-                self.inter_AB = Inter(in_ch=encoder_out_ch, lowest_dense_res=lowest_dense_res, ae_ch=ae_dims, ae_out_ch=ae_dims*2, name='inter_AB')
+                self.inter_AB = Inter(in_ch=encoder_out_ch, lowest_dense_res=lowest_dense_res, ae_ch=ae_dims, ae_out_ch=ae_dims*2, is_hd=is_hd, name='inter_AB')
-                self.inter_B  = Inter(in_ch=encoder_out_ch, lowest_dense_res=lowest_dense_res, ae_ch=ae_dims, ae_out_ch=ae_dims*2, name='inter_B')
+                self.inter_B  = Inter(in_ch=encoder_out_ch, lowest_dense_res=lowest_dense_res, ae_ch=ae_dims, ae_out_ch=ae_dims*2, is_hd=is_hd, name='inter_B')
 
                 inter_AB_out_ch = self.inter_AB.compute_output_channels ( (nn.tf_floatx, (None,encoder_out_ch)))
                 inter_B_out_ch = self.inter_B.compute_output_channels ( (nn.tf_floatx, (None,encoder_out_ch)))
                 inters_out_ch = inter_AB_out_ch+inter_B_out_ch
-                self.decoder = Decoder(in_ch=inters_out_ch, d_ch=d_dims, d_mask_ch=d_mask_dims, is_hd='hd' in archi, name='decoder')
+                self.decoder = Decoder(in_ch=inters_out_ch, d_ch=d_dims, d_mask_ch=d_mask_dims, is_hd=is_hd, name='decoder')
 
                 self.model_filename_list += [ [self.encoder,  'encoder.npy'],
                                               [self.inter_AB, 'inter_AB.npy'],
@@ -439,7 +227,7 @@ class SAEHDModel(ModelBase):
 
                 # Initialize optimizers
                 lr=5e-5
-                lr_dropout = 0.3 if self.options['lr_dropout'] else 1.0
+                lr_dropout = 0.3 if self.options['lr_dropout'] and not self.pretrain else 1.0
                 clipnorm = 1.0 if self.options['clipgrad'] else 0.0
                 self.src_dst_opt = nn.TFRMSpropOptimizer(lr=lr, lr_dropout=lr_dropout, clipnorm=clipnorm, name='src_dst_opt')
                 self.model_filename_list += [ (self.src_dst_opt, 'src_dst_opt.npy') ]