mirror of
https://github.com/iperov/DeepFaceLab.git
synced 2024-12-25 23:41:12 -08:00
195 lines
7.4 KiB
Python
195 lines
7.4 KiB
Python
import numpy as np
|
|
from core.leras import nn
|
|
tf = nn.tf
|
|
|
|
patch_discriminator_kernels = \
|
|
{ 1 : (512, [ [1,1] ]),
|
|
2 : (512, [ [2,1] ]),
|
|
3 : (512, [ [2,1], [2,1] ]),
|
|
4 : (512, [ [2,2], [2,2] ]),
|
|
5 : (512, [ [3,2], [2,2] ]),
|
|
6 : (512, [ [4,2], [2,2] ]),
|
|
7 : (512, [ [3,2], [3,2] ]),
|
|
8 : (512, [ [4,2], [3,2] ]),
|
|
9 : (512, [ [3,2], [4,2] ]),
|
|
10 : (512, [ [4,2], [4,2] ]),
|
|
11 : (512, [ [3,2], [3,2], [2,1] ]),
|
|
12 : (512, [ [4,2], [3,2], [2,1] ]),
|
|
13 : (512, [ [3,2], [4,2], [2,1] ]),
|
|
14 : (512, [ [4,2], [4,2], [2,1] ]),
|
|
15 : (512, [ [3,2], [3,2], [3,1] ]),
|
|
16 : (512, [ [4,2], [3,2], [3,1] ]),
|
|
17 : (512, [ [3,2], [4,2], [3,1] ]),
|
|
18 : (512, [ [4,2], [4,2], [3,1] ]),
|
|
19 : (512, [ [3,2], [3,2], [4,1] ]),
|
|
20 : (512, [ [4,2], [3,2], [4,1] ]),
|
|
21 : (512, [ [3,2], [4,2], [4,1] ]),
|
|
22 : (512, [ [4,2], [4,2], [4,1] ]),
|
|
23 : (256, [ [3,2], [3,2], [3,2], [2,1] ]),
|
|
24 : (256, [ [4,2], [3,2], [3,2], [2,1] ]),
|
|
25 : (256, [ [3,2], [4,2], [3,2], [2,1] ]),
|
|
26 : (256, [ [4,2], [4,2], [3,2], [2,1] ]),
|
|
27 : (256, [ [3,2], [4,2], [4,2], [2,1] ]),
|
|
28 : (256, [ [4,2], [3,2], [4,2], [2,1] ]),
|
|
29 : (256, [ [3,2], [4,2], [4,2], [2,1] ]),
|
|
30 : (256, [ [4,2], [4,2], [4,2], [2,1] ]),
|
|
31 : (256, [ [3,2], [3,2], [3,2], [3,1] ]),
|
|
32 : (256, [ [4,2], [3,2], [3,2], [3,1] ]),
|
|
33 : (256, [ [3,2], [4,2], [3,2], [3,1] ]),
|
|
34 : (256, [ [4,2], [4,2], [3,2], [3,1] ]),
|
|
35 : (256, [ [3,2], [4,2], [4,2], [3,1] ]),
|
|
36 : (256, [ [4,2], [3,2], [4,2], [3,1] ]),
|
|
37 : (256, [ [3,2], [4,2], [4,2], [3,1] ]),
|
|
38 : (256, [ [4,2], [4,2], [4,2], [3,1] ]),
|
|
39 : (256, [ [3,2], [3,2], [3,2], [4,1] ]),
|
|
40 : (256, [ [4,2], [3,2], [3,2], [4,1] ]),
|
|
41 : (256, [ [3,2], [4,2], [3,2], [4,1] ]),
|
|
42 : (256, [ [4,2], [4,2], [3,2], [4,1] ]),
|
|
43 : (256, [ [3,2], [4,2], [4,2], [4,1] ]),
|
|
44 : (256, [ [4,2], [3,2], [4,2], [4,1] ]),
|
|
45 : (256, [ [3,2], [4,2], [4,2], [4,1] ]),
|
|
46 : (256, [ [4,2], [4,2], [4,2], [4,1] ]),
|
|
}
|
|
|
|
|
|
class PatchDiscriminator(nn.ModelBase):
|
|
def on_build(self, patch_size, in_ch, base_ch=None, conv_kernel_initializer=None):
|
|
suggested_base_ch, kernels_strides = patch_discriminator_kernels[patch_size]
|
|
|
|
if base_ch is None:
|
|
base_ch = suggested_base_ch
|
|
|
|
prev_ch = in_ch
|
|
self.convs = []
|
|
for i, (kernel_size, strides) in enumerate(kernels_strides):
|
|
cur_ch = base_ch * min( (2**i), 8 )
|
|
|
|
self.convs.append ( nn.Conv2D( prev_ch, cur_ch, kernel_size=kernel_size, strides=strides, 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.PatchDiscriminator = PatchDiscriminator
|
|
|
|
class UNetPatchDiscriminator(nn.ModelBase):
|
|
"""
|
|
Inspired by https://arxiv.org/abs/2002.12655 "A U-Net Based Discriminator for Generative Adversarial Networks"
|
|
"""
|
|
def calc_receptive_field_size(self, layers):
|
|
"""
|
|
result the same as https://fomoro.com/research/article/receptive-field-calculatorindex.html
|
|
"""
|
|
rf = 0
|
|
ts = 1
|
|
for i, (k, s) in enumerate(layers):
|
|
if i == 0:
|
|
rf = k
|
|
else:
|
|
rf += (k-1)*ts
|
|
ts *= s
|
|
return rf
|
|
|
|
def find_archi(self, target_patch_size, max_layers=9):
|
|
"""
|
|
Find the best configuration of layers using only 3x3 convs for target patch size
|
|
"""
|
|
s = {}
|
|
for layers_count in range(1,max_layers+1):
|
|
val = 1 << (layers_count-1)
|
|
while True:
|
|
val -= 1
|
|
|
|
layers = []
|
|
sum_st = 0
|
|
layers.append ( [3, 2])
|
|
sum_st += 2
|
|
for i in range(layers_count-1):
|
|
st = 1 + (1 if val & (1 << i) !=0 else 0 )
|
|
layers.append ( [3, st ])
|
|
sum_st += st
|
|
|
|
rf = self.calc_receptive_field_size(layers)
|
|
|
|
s_rf = s.get(rf, None)
|
|
if s_rf is None:
|
|
s[rf] = (layers_count, sum_st, layers)
|
|
else:
|
|
if layers_count < s_rf[0] or \
|
|
( layers_count == s_rf[0] and sum_st > s_rf[1] ):
|
|
s[rf] = (layers_count, sum_st, layers)
|
|
|
|
if val == 0:
|
|
break
|
|
|
|
x = sorted(list(s.keys()))
|
|
q=x[np.abs(np.array(x)-target_patch_size).argmin()]
|
|
return s[q][2]
|
|
|
|
def on_build(self, patch_size, in_ch, base_ch = 16, use_fp16 = False):
|
|
self.use_fp16 = use_fp16
|
|
conv_dtype = tf.float16 if use_fp16 else tf.float32
|
|
|
|
class ResidualBlock(nn.ModelBase):
|
|
def on_build(self, ch, kernel_size=3 ):
|
|
self.conv1 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME', dtype=conv_dtype)
|
|
self.conv2 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME', dtype=conv_dtype)
|
|
|
|
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
|
|
|
|
prev_ch = in_ch
|
|
self.convs = []
|
|
self.upconvs = []
|
|
layers = self.find_archi(patch_size)
|
|
|
|
level_chs = { i-1:v for i,v in enumerate([ min( base_ch * (2**i), 512 ) for i in range(len(layers)+1)]) }
|
|
|
|
self.in_conv = nn.Conv2D( in_ch, level_chs[-1], kernel_size=1, padding='VALID', dtype=conv_dtype)
|
|
|
|
for i, (kernel_size, strides) in enumerate(layers):
|
|
self.convs.append ( nn.Conv2D( level_chs[i-1], level_chs[i], kernel_size=kernel_size, strides=strides, padding='SAME', dtype=conv_dtype) )
|
|
|
|
self.upconvs.insert (0, nn.Conv2DTranspose( level_chs[i]*(2 if i != len(layers)-1 else 1), level_chs[i-1], kernel_size=kernel_size, strides=strides, padding='SAME', dtype=conv_dtype) )
|
|
|
|
self.out_conv = nn.Conv2D( level_chs[-1]*2, 1, kernel_size=1, padding='VALID', dtype=conv_dtype)
|
|
|
|
self.center_out = nn.Conv2D( level_chs[len(layers)-1], 1, kernel_size=1, padding='VALID', dtype=conv_dtype)
|
|
self.center_conv = nn.Conv2D( level_chs[len(layers)-1], level_chs[len(layers)-1], kernel_size=1, padding='VALID', dtype=conv_dtype)
|
|
|
|
|
|
def forward(self, x):
|
|
if self.use_fp16:
|
|
x = tf.cast(x, tf.float16)
|
|
|
|
x = tf.nn.leaky_relu( self.in_conv(x), 0.2 )
|
|
|
|
encs = []
|
|
for conv in self.convs:
|
|
encs.insert(0, x)
|
|
x = tf.nn.leaky_relu( conv(x), 0.2 )
|
|
|
|
center_out, x = self.center_out(x), tf.nn.leaky_relu( self.center_conv(x), 0.2 )
|
|
|
|
for i, (upconv, enc) in enumerate(zip(self.upconvs, encs)):
|
|
x = tf.nn.leaky_relu( upconv(x), 0.2 )
|
|
x = tf.concat( [enc, x], axis=nn.conv2d_ch_axis)
|
|
|
|
x = self.out_conv(x)
|
|
|
|
if self.use_fp16:
|
|
center_out = tf.cast(center_out, tf.float32)
|
|
x = tf.cast(x, tf.float32)
|
|
|
|
return center_out, x
|
|
|
|
nn.UNetPatchDiscriminator = UNetPatchDiscriminator
|