mirror of
https://github.com/iperov/DeepFaceLab.git
synced 2025-07-06 04:52:13 -07:00
ada60ccefe
lr_dropout is now disabled in pretraining mode changed help message for lr_dropout and random_warp
565 lines
No EOL
27 KiB
Python
565 lines
No EOL
27 KiB
Python
|
|
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 |