github/DeepFaceLab
1
0
Fork 0
mirror of https://github.com/iperov/DeepFaceLab.git synced 2025-07-06 13:02:15 -07:00
Code Issues Projects Releases Packages Wiki Activity Actions

Upgraded to TF version 1.13.2

Browse source 
Removed the wait at first launch for most graphics cards.

Increased speed of training by 10-20%, but you have to retrain all models from scratch.

SAEHD:

added option 'use float16'
	Experimental option. Reduces the model size by half.
	Increases the speed of training.
	Decreases the accuracy of the model.
	The model may collapse or not train.
	Model may not learn the mask in large resolutions.

true_face_training option is replaced by
"True face power". 0.0000 .. 1.0
Experimental option. Discriminates the result face to be more like the src face. Higher value - stronger discrimination.
Comparison - https://i.imgur.com/czScS9q.png
This commit is contained in:
Colombo 2020-01-25 21:58:19 +04:00
parent a3dfcb91b9
commit 76ca79216e
49 changed files with 1320 additions and 1297 deletions
Download patch file Download diff file Expand all files Collapse all files

5
core/leras/initializers.py
View file

@ -11,11 +11,8 @@ def initialize_initializers(nn):
class initializers(): class initializers():
class ca (init_ops.Initializer): class ca (init_ops.Initializer):
def __init__(self, dtype=None):
pass
def __call__(self, shape, dtype=None, partition_info=None): def __call__(self, shape, dtype=None, partition_info=None):
return tf.zeros( shape, name="_cai_") return tf.zeros( shape, dtype=dtype, name="_cai_")
@staticmethod @staticmethod
def generate_batch( data_list, eps_std=0.05 ): def generate_batch( data_list, eps_std=0.05 ):

137
core/leras/layers.py
View file

@ -65,6 +65,8 @@ def initialize_layers(nn):
sub_w_name = "/".join(w_name_split[1:]) sub_w_name = "/".join(w_name_split[1:])
w_val = d.get(sub_w_name, None) w_val = d.get(sub_w_name, None)
w_val = np.reshape( w_val, w.shape.as_list() )
if w_val is None: if w_val is None:
io.log_err(f"Weight {w.name} was not loaded from file {filename}") io.log_err(f"Weight {w.name} was not loaded from file {filename}")
tuples.append ( (w, w.initializer) ) tuples.append ( (w, w.initializer) )
@ -255,13 +257,22 @@ def initialize_layers(nn):
return result_shapes[0] if not_list else result_shapes return result_shapes[0] if not_list else result_shapes
def compute_output_channels(self, shapes):
shape = self.compute_output_shape(shapes)
shape_len = len(shape)
if shape_len == 4:
if nn.data_format == "NCHW":
return shape[1]
return shape[-1]
def build_for_run(self, shapes_list): def build_for_run(self, shapes_list):
if not isinstance(shapes_list, list): if not isinstance(shapes_list, list):
raise ValueError("shapes_list must be a list.") raise ValueError("shapes_list must be a list.")
self.run_placeholders = [] self.run_placeholders = []
for dtype,sh in shapes_list: for dtype,sh in shapes_list:
self.run_placeholders.append ( tf.placeholder(dtype, (None,)+sh) ) self.run_placeholders.append ( tf.placeholder(dtype, sh) )
self.run_output = self.__call__(self.run_placeholders) self.run_output = self.__call__(self.run_placeholders)
@ -292,6 +303,9 @@ def initialize_layers(nn):
if not isinstance(dilations, int): if not isinstance(dilations, int):
raise ValueError ("dilations must be an int type") raise ValueError ("dilations must be an int type")
if dtype is None:
dtype = nn.tf_floatx
if isinstance(padding, str): if isinstance(padding, str):
if padding == "SAME": if padding == "SAME":
padding = ( (kernel_size - 1) * dilations + 1 ) // 2 padding = ( (kernel_size - 1) * dilations + 1 ) // 2
@ -302,36 +316,47 @@ def initialize_layers(nn):
if isinstance(padding, int): if isinstance(padding, int):
if padding != 0: if padding != 0:
if nn.data_format == "NHWC":
padding = [ [0,0], [padding,padding], [padding,padding], [0,0] ] padding = [ [0,0], [padding,padding], [padding,padding], [0,0] ]
else:
padding = [ [0,0], [0,0], [padding,padding], [padding,padding] ]
else: else:
padding = None padding = None
if nn.data_format == "NHWC":
strides = [1,strides,strides,1]
else:
strides = [1,1,strides,strides]
if nn.data_format == "NHWC":
dilations = [1,dilations,dilations,1]
else:
dilations = [1,1,dilations,dilations]
self.in_ch = in_ch self.in_ch = in_ch
self.out_ch = out_ch self.out_ch = out_ch
self.kernel_size = kernel_size self.kernel_size = kernel_size
self.strides = [1,strides,strides,1] self.strides = strides
self.padding = padding self.padding = padding
self.dilations = [1,dilations,dilations,1] self.dilations = dilations
self.use_bias = use_bias self.use_bias = use_bias
self.use_wscale = use_wscale self.use_wscale = use_wscale
self.kernel_initializer = None if use_wscale else kernel_initializer self.kernel_initializer = kernel_initializer
self.bias_initializer = bias_initializer self.bias_initializer = bias_initializer
self.trainable = trainable self.trainable = trainable
if dtype is None:
dtype = nn.tf_floatx
self.dtype = dtype self.dtype = dtype
super().__init__(**kwargs) super().__init__(**kwargs)
def build_weights(self): def build_weights(self):
kernel_initializer = self.kernel_initializer kernel_initializer = self.kernel_initializer
if kernel_initializer is None:
if self.use_wscale: if self.use_wscale:
gain = 1.0 if self.kernel_size == 1 else np.sqrt(2) gain = 1.0 if self.kernel_size == 1 else np.sqrt(2)
fan_in = self.kernel_size*self.kernel_size*self.in_ch fan_in = self.kernel_size*self.kernel_size*self.in_ch
he_std = gain / np.sqrt(fan_in) # He init he_std = gain / np.sqrt(fan_in) # He init
self.wscale = tf.constant(he_std, dtype=self.dtype ) self.wscale = tf.constant(he_std, dtype=self.dtype )
kernel_initializer = tf.initializers.random_normal(0, 1.0, dtype=self.dtype) kernel_initializer = tf.initializers.random_normal(0, 1.0, dtype=self.dtype)
else:
if kernel_initializer is None:
kernel_initializer = tf.initializers.glorot_uniform(dtype=self.dtype) kernel_initializer = tf.initializers.glorot_uniform(dtype=self.dtype)
self.weight = tf.get_variable("weight", (self.kernel_size,self.kernel_size,self.in_ch,self.out_ch), dtype=self.dtype, initializer=kernel_initializer, trainable=self.trainable ) self.weight = tf.get_variable("weight", (self.kernel_size,self.kernel_size,self.in_ch,self.out_ch), dtype=self.dtype, initializer=kernel_initializer, trainable=self.trainable )
@ -341,7 +366,7 @@ def initialize_layers(nn):
if bias_initializer is None: if bias_initializer is None:
bias_initializer = tf.initializers.zeros(dtype=self.dtype) bias_initializer = tf.initializers.zeros(dtype=self.dtype)
self.bias = tf.get_variable("bias", (1,1,1,self.out_ch), dtype=self.dtype, initializer=bias_initializer, trainable=self.trainable ) self.bias = tf.get_variable("bias", (self.out_ch,), dtype=self.dtype, initializer=bias_initializer, trainable=self.trainable )
def get_weights(self): def get_weights(self):
weights = [self.weight] weights = [self.weight]
@ -357,9 +382,13 @@ def initialize_layers(nn):
if self.padding is not None: if self.padding is not None:
x = tf.pad (x, self.padding, mode='CONSTANT') x = tf.pad (x, self.padding, mode='CONSTANT')
x = tf.nn.conv2d(x, weight, self.strides, 'VALID', dilations=self.dilations) x = tf.nn.conv2d(x, weight, self.strides, 'VALID', dilations=self.dilations, data_format=nn.data_format)
if self.use_bias: if self.use_bias:
x = x + self.bias if nn.data_format == "NHWC":
bias = tf.reshape (self.bias, (1,1,1,self.out_ch) )
else:
bias = tf.reshape (self.bias, (1,self.out_ch,1,1) )
x = tf.add(x, bias)
return x return x
def __str__(self): def __str__(self):
@ -376,6 +405,10 @@ def initialize_layers(nn):
def __init__(self, in_ch, out_ch, kernel_size, strides=2, padding='SAME', use_bias=True, use_wscale=False, kernel_initializer=None, bias_initializer=None, trainable=True, dtype=None, **kwargs ): def __init__(self, in_ch, out_ch, kernel_size, strides=2, padding='SAME', use_bias=True, use_wscale=False, kernel_initializer=None, bias_initializer=None, trainable=True, dtype=None, **kwargs ):
if not isinstance(strides, int): if not isinstance(strides, int):
raise ValueError ("strides must be an int type") raise ValueError ("strides must be an int type")
if dtype is None:
dtype = nn.tf_floatx
self.in_ch = in_ch self.in_ch = in_ch
self.out_ch = out_ch self.out_ch = out_ch
self.kernel_size = kernel_size self.kernel_size = kernel_size
@ -383,33 +416,30 @@ def initialize_layers(nn):
self.padding = padding self.padding = padding
self.use_bias = use_bias self.use_bias = use_bias
self.use_wscale = use_wscale self.use_wscale = use_wscale
self.kernel_initializer = None if use_wscale else kernel_initializer self.kernel_initializer = kernel_initializer
self.bias_initializer = bias_initializer self.bias_initializer = bias_initializer
self.trainable = trainable self.trainable = trainable
if dtype is None:
dtype = nn.tf_floatx
self.dtype = dtype self.dtype = dtype
super().__init__(**kwargs) super().__init__(**kwargs)
def build_weights(self): def build_weights(self):
kernel_initializer = self.kernel_initializer kernel_initializer = self.kernel_initializer
if kernel_initializer is None:
if self.use_wscale: if self.use_wscale:
gain = 1.0 if self.kernel_size == 1 else np.sqrt(2) gain = 1.0 if self.kernel_size == 1 else np.sqrt(2)
fan_in = self.kernel_size*self.kernel_size*self.in_ch fan_in = self.kernel_size*self.kernel_size*self.in_ch
he_std = gain / np.sqrt(fan_in) # He init he_std = gain / np.sqrt(fan_in) # He init
self.wscale = tf.constant(he_std, dtype=self.dtype ) self.wscale = tf.constant(he_std, dtype=self.dtype )
kernel_initializer = tf.initializers.random_normal(0, 1.0, dtype=self.dtype) kernel_initializer = tf.initializers.random_normal(0, 1.0, dtype=self.dtype)
else: if kernel_initializer is None:
kernel_initializer = tf.initializers.glorot_uniform(dtype=self.dtype) kernel_initializer = tf.initializers.glorot_uniform(dtype=self.dtype)
self.weight = tf.get_variable("weight", (self.kernel_size,self.kernel_size,self.out_ch,self.in_ch), dtype=self.dtype, initializer=kernel_initializer, trainable=self.trainable ) self.weight = tf.get_variable("weight", (self.kernel_size,self.kernel_size,self.out_ch,self.in_ch), dtype=self.dtype, initializer=kernel_initializer, trainable=self.trainable )
if self.use_bias: if self.use_bias:
bias_initializer = self.bias_initializer bias_initializer = self.bias_initializer
if bias_initializer is None: if bias_initializer is None:
bias_initializer = tf.initializers.zeros(dtype=self.dtype) bias_initializer = tf.initializers.zeros(dtype=self.dtype)
self.bias = tf.get_variable("bias", (1,1,1,self.out_ch), dtype=self.dtype, initializer=bias_initializer, trainable=self.trainable )
self.bias = tf.get_variable("bias", (self.out_ch,), dtype=self.dtype, initializer=bias_initializer, trainable=self.trainable )
def get_weights(self): def get_weights(self):
weights = [self.weight] weights = [self.weight]
@ -420,21 +450,34 @@ def initialize_layers(nn):
def __call__(self, x): def __call__(self, x):
shape = x.shape shape = x.shape
if nn.data_format == "NHWC":
h,w,c = shape[1], shape[2], shape[3] h,w,c = shape[1], shape[2], shape[3]
output_shape = tf.stack ( (tf.shape(x)[0], output_shape = tf.stack ( (tf.shape(x)[0],
self.deconv_length(w, self.strides, self.kernel_size, self.padding), self.deconv_length(w, self.strides, self.kernel_size, self.padding),
self.deconv_length(h, self.strides, self.kernel_size, self.padding), self.deconv_length(h, self.strides, self.kernel_size, self.padding),
self.out_ch) ) self.out_ch) )
strides = [1,self.strides,self.strides,1]
else:
c,h,w = shape[1], shape[2], shape[3]
output_shape = tf.stack ( (tf.shape(x)[0],
self.out_ch,
self.deconv_length(w, self.strides, self.kernel_size, self.padding),
self.deconv_length(h, self.strides, self.kernel_size, self.padding),
) )
strides = [1,1,self.strides,self.strides]
weight = self.weight weight = self.weight
if self.use_wscale: if self.use_wscale:
weight = weight * self.wscale weight = weight * self.wscale
x = tf.nn.conv2d_transpose(x, weight, output_shape, [1,self.strides,self.strides,1], padding=self.padding) x = tf.nn.conv2d_transpose(x, weight, output_shape, strides, padding=self.padding, data_format=nn.data_format)
if self.use_bias: if self.use_bias:
x = x + self.bias if nn.data_format == "NHWC":
bias = tf.reshape (self.bias, (1,1,1,self.out_ch) )
else:
bias = tf.reshape (self.bias, (1,self.out_ch,1,1) )
x = tf.add(x, bias)
return x return x
def __str__(self): def __str__(self):
@ -459,10 +502,13 @@ def initialize_layers(nn):
def __init__(self, filt_size=3, stride=2, **kwargs ): def __init__(self, filt_size=3, stride=2, **kwargs ):
self.strides = [1,stride,stride,1] self.strides = [1,stride,stride,1]
self.filt_size = filt_size self.filt_size = filt_size
self.padding = [ [0,0], pad = [ int(1.*(filt_size-1)/2), int(np.ceil(1.*(filt_size-1)/2)) ]
[ int(1.*(filt_size-1)/2), int(np.ceil(1.*(filt_size-1)/2)) ],
[ int(1.*(filt_size-1)/2), int(np.ceil(1.*(filt_size-1)/2)) ], if nn.data_format == "NHWC":
[0,0] ] self.padding = [ [0,0], pad, pad, [0,0] ]
else:
self.padding = [ [0,0], [0,0], pad, pad ]
if(self.filt_size==1): if(self.filt_size==1):
a = np.array([1.,]) a = np.array([1.,])
elif(self.filt_size==2): elif(self.filt_size==2):
@ -512,7 +558,8 @@ def initialize_layers(nn):
self.bias_initializer = bias_initializer self.bias_initializer = bias_initializer
self.trainable = trainable self.trainable = trainable
if dtype is None: if dtype is None:
dtype = tf.float32 dtype = nn.tf_floatx
self.dtype = dtype self.dtype = dtype
super().__init__(**kwargs) super().__init__(**kwargs)
@ -523,14 +570,15 @@ def initialize_layers(nn):
weight_shape = (self.in_ch,self.out_ch) weight_shape = (self.in_ch,self.out_ch)
kernel_initializer = self.kernel_initializer kernel_initializer = self.kernel_initializer
if kernel_initializer is None:
if self.use_wscale: if self.use_wscale:
gain = 1.0 gain = 1.0
fan_in = np.prod( weight_shape[:-1] ) fan_in = np.prod( weight_shape[:-1] )
he_std = gain / np.sqrt(fan_in) # He init he_std = gain / np.sqrt(fan_in) # He init
self.wscale = tf.constant(he_std, dtype=self.dtype ) self.wscale = tf.constant(he_std, dtype=self.dtype )
kernel_initializer = tf.initializers.random_normal(0, 1.0, dtype=self.dtype) kernel_initializer = tf.initializers.random_normal(0, 1.0, dtype=self.dtype)
else:
if kernel_initializer is None:
kernel_initializer = tf.initializers.glorot_uniform(dtype=self.dtype) kernel_initializer = tf.initializers.glorot_uniform(dtype=self.dtype)
self.weight = tf.get_variable("weight", weight_shape, dtype=self.dtype, initializer=kernel_initializer, trainable=self.trainable ) self.weight = tf.get_variable("weight", weight_shape, dtype=self.dtype, initializer=kernel_initializer, trainable=self.trainable )
@ -539,7 +587,7 @@ def initialize_layers(nn):
bias_initializer = self.bias_initializer bias_initializer = self.bias_initializer
if bias_initializer is None: if bias_initializer is None:
bias_initializer = tf.initializers.zeros(dtype=self.dtype) bias_initializer = tf.initializers.zeros(dtype=self.dtype)
self.bias = tf.get_variable("bias", (1,self.out_ch), dtype=self.dtype, initializer=bias_initializer, trainable=self.trainable ) self.bias = tf.get_variable("bias", (self.out_ch,), dtype=self.dtype, initializer=bias_initializer, trainable=self.trainable )
def get_weights(self): def get_weights(self):
weights = [self.weight] weights = [self.weight]
@ -559,7 +607,7 @@ def initialize_layers(nn):
x = tf.reduce_max(x, axis=-1) x = tf.reduce_max(x, axis=-1)
if self.use_bias: if self.use_bias:
x = x + self.bias x = tf.add(x, tf.reshape(self.bias, (1,self.out_ch) ) )
return x return x
nn.Dense = Dense nn.Dense = Dense
@ -568,31 +616,38 @@ def initialize_layers(nn):
""" """
currently not for training currently not for training
""" """
def __init__(self, dim, eps=1e-05, momentum=0.1, dtype=None, **kwargs ): def __init__(self, dim, eps=1e-05, momentum=0.1, dtype=None, **kwargs):
self.dim = dim self.dim = dim
self.eps = eps self.eps = eps
self.momentum = momentum self.momentum = momentum
if dtype is None: if dtype is None:
dtype = nn.tf_floatx dtype = nn.tf_floatx
self.dtype = dtype self.dtype = dtype
self.shape = (1,1,1,dim)
super().__init__(**kwargs) super().__init__(**kwargs)
def build_weights(self): def build_weights(self):
self.weight = tf.get_variable("weight", self.shape, dtype=self.dtype, initializer=tf.initializers.ones() ) self.weight = tf.get_variable("weight", (self.dim,), dtype=self.dtype, initializer=tf.initializers.ones() )
self.bias = tf.get_variable("bias", self.shape, dtype=self.dtype, initializer=tf.initializers.zeros() ) self.bias = tf.get_variable("bias", (self.dim,), dtype=self.dtype, initializer=tf.initializers.zeros() )
self.running_mean = tf.get_variable("running_mean", self.shape, dtype=self.dtype, initializer=tf.initializers.zeros(), trainable=False ) self.running_mean = tf.get_variable("running_mean", (self.dim,), dtype=self.dtype, initializer=tf.initializers.zeros(), trainable=False )
self.running_var = tf.get_variable("running_var", self.shape, dtype=self.dtype, initializer=tf.initializers.zeros(), trainable=False ) self.running_var = tf.get_variable("running_var", (self.dim,), dtype=self.dtype, initializer=tf.initializers.zeros(), trainable=False )
def get_weights(self): def get_weights(self):
return [self.weight, self.bias, self.running_mean, self.running_var] return [self.weight, self.bias, self.running_mean, self.running_var]
def __call__(self, x): def __call__(self, x):
x = (x - self.running_mean) / tf.sqrt( self.running_var + self.eps ) if nn.data_format == "NHWC":
x *= self.weight shape = (1,1,1,self.dim)
x += self.bias else:
shape = (1,self.dim,1,1)
weight = tf.reshape ( self.weight , shape )
bias = tf.reshape ( self.bias , shape )
running_mean = tf.reshape ( self.running_mean, shape )
running_var = tf.reshape ( self.running_var , shape )
x = (x - running_mean) / tf.sqrt( running_var + self.eps )
x *= weight
x += bias
return x return x
nn.BatchNorm2D = BatchNorm2D nn.BatchNorm2D = BatchNorm2D

106
core/leras/nn.py
View file

@ -13,20 +13,35 @@ Provides:
Reasons why we cannot import tensorflow or any tensorflow.sub modules right here: Reasons why we cannot import tensorflow or any tensorflow.sub modules right here:
1) change env variables based on DeviceConfig before import tensorflow 1) change env variables based on DeviceConfig before import tensorflow
2) multiprocesses will import tensorflow every spawn 2) multiprocesses will import tensorflow every spawn
NCHW speed up training for 10-20%.
""" """
import os import os
import sys import sys
from pathlib import Path from pathlib import Path
import numpy as np
from core.interact import interact as io from core.interact import interact as io
from .device import Devices from .device import Devices
class nn(): class nn():
current_DeviceConfig = None current_DeviceConfig = None
tf = None tf = None
tf_sess = None tf_sess = None
tf_sess_config = None tf_sess_config = None
tf_default_device = None
data_format = None
conv2d_ch_axis = None
conv2d_spatial_axes = None
tf_floatx = None
np_floatx = None
# Tensor ops # Tensor ops
tf_get_value = None tf_get_value = None
@ -34,17 +49,18 @@ class nn():
tf_gradients = None tf_gradients = None
tf_average_gv_list = None tf_average_gv_list = None
tf_average_tensor_list = None tf_average_tensor_list = None
tf_dot = None tf_concat = None
tf_gelu = None tf_gelu = None
tf_upsample2d = None tf_upsample2d = None
tf_upsample2d_bilinear = None tf_upsample2d_bilinear = None
tf_flatten = None tf_flatten = None
tf_reshape_4D = None
tf_random_binomial = None tf_random_binomial = None
tf_gaussian_blur = None tf_gaussian_blur = None
tf_style_loss = None tf_style_loss = None
tf_channel_histogram = None
tf_histogram = None
tf_dssim = None tf_dssim = None
tf_space_to_depth = None
tf_depth_to_space = None
# Layers # Layers
Saveable = None Saveable = None
@ -64,7 +80,8 @@ class nn():
TFRMSpropOptimizer = None TFRMSpropOptimizer = None
@staticmethod @staticmethod
def initialize(device_config=None): def initialize(device_config=None, floatx="float32", data_format="NHWC"):
if nn.tf is None: if nn.tf is None:
if device_config is None: if device_config is None:
device_config = nn.getCurrentDeviceConfig() device_config = nn.getCurrentDeviceConfig()
@ -74,11 +91,8 @@ class nn():
if 'CUDA_VISIBLE_DEVICES' in os.environ.keys(): if 'CUDA_VISIBLE_DEVICES' in os.environ.keys():
os.environ.pop('CUDA_VISIBLE_DEVICES') os.environ.pop('CUDA_VISIBLE_DEVICES')
os.environ['CUDA_​CACHE_​MAXSIZE'] = '536870912' #512Mb (32mb default)
first_run = False first_run = False
if len(device_config.devices) != 0:
if not device_config.cpu_only:
if sys.platform[0:3] == 'win': if sys.platform[0:3] == 'win':
if all( [ x.name == device_config.devices[0].name for x in device_config.devices ] ): if all( [ x.name == device_config.devices[0].name for x in device_config.devices ] ):
devices_str = "_" + device_config.devices[0].name.replace(' ','_') devices_str = "_" + device_config.devices[0].name.replace(' ','_')
@ -92,6 +106,7 @@ class nn():
first_run = True first_run = True
os.environ['CUDA_CACHE_PATH'] = str(compute_cache_path) os.environ['CUDA_CACHE_PATH'] = str(compute_cache_path)
os.environ['CUDA_​CACHE_​MAXSIZE'] = '536870912' #512Mb (32mb default)
os.environ['TF_MIN_GPU_MULTIPROCESSOR_COUNT'] = '2' os.environ['TF_MIN_GPU_MULTIPROCESSOR_COUNT'] = '2'
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' # tf log errors only os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' # tf log errors only
@ -102,11 +117,16 @@ class nn():
io.log_info("Caching GPU kernels...") io.log_info("Caching GPU kernels...")
import tensorflow as tf import tensorflow as tf
import logging
logging.getLogger('tensorflow').setLevel(logging.ERROR)
nn.tf = tf nn.tf = tf
if device_config.cpu_only: if len(device_config.devices) == 0:
nn.tf_default_device = "/CPU:0"
config = tf.ConfigProto(device_count={'GPU': 0}) config = tf.ConfigProto(device_count={'GPU': 0})
else: else:
nn.tf_default_device = "/GPU:0"
config = tf.ConfigProto() config = tf.ConfigProto()
config.gpu_options.visible_device_list = ','.join([str(device.index) for device in device_config.devices]) config.gpu_options.visible_device_list = ','.join([str(device.index) for device in device_config.devices])
@ -114,9 +134,6 @@ class nn():
config.gpu_options.allow_growth = True config.gpu_options.allow_growth = True
nn.tf_sess_config = config nn.tf_sess_config = config
nn.tf_floatx = nn.tf.float32 #nn.tf.float16 if device_config.use_fp16 else nn.tf.float32
nn.np_floatx = nn.tf_floatx.as_numpy_dtype
from .tensor_ops import initialize_tensor_ops from .tensor_ops import initialize_tensor_ops
from .layers import initialize_layers from .layers import initialize_layers
from .initializers import initialize_initializers from .initializers import initialize_initializers
@ -130,10 +147,68 @@ class nn():
if nn.tf_sess is None: if nn.tf_sess is None:
nn.tf_sess = tf.Session(config=nn.tf_sess_config) nn.tf_sess = tf.Session(config=nn.tf_sess_config)
if floatx == "float32":
floatx = nn.tf.float32
elif floatx == "float16":
floatx = nn.tf.float16
else:
raise ValueError(f"unsupported floatx {floatx}")
nn.set_floatx(floatx)
nn.set_data_format(data_format)
@staticmethod @staticmethod
def initialize_main_env(): def initialize_main_env():
Devices.initialize_main_env() Devices.initialize_main_env()
@staticmethod
def set_floatx(tf_dtype):
"""
set default float type for all layers when dtype is None for them
"""
nn.tf_floatx = tf_dtype
nn.np_floatx = tf_dtype.as_numpy_dtype
@staticmethod
def set_data_format(data_format):
if data_format != "NHWC" and data_format != "NCHW":
raise ValueError(f"unsupported data_format {data_format}")
nn.data_format = data_format
if data_format == "NHWC":
nn.conv2d_ch_axis = 3
nn.conv2d_spatial_axes = [1,2]
elif data_format == "NCHW":
nn.conv2d_ch_axis = 1
nn.conv2d_spatial_axes = [2,3]
@staticmethod
def get4Dshape ( w, h, c, data_format=None ):
"""
returns 4D shape based on current data_format
"""
if data_format is None:
data_format = nn.data_format
if data_format == "NHWC":
return (None,h,w,c)
else:
return (None,c,h,w)
@staticmethod
def to_data_format( x, to_data_format, from_data_format=None):
if from_data_format is None:
from_data_format = nn.data_format
if to_data_format == from_data_format:
return x
if to_data_format == "NHWC":
return np.transpose(x, (0,2,3,1) )
elif to_data_format == "NCHW":
return np.transpose(x, (0,3,1,2) )
else:
raise ValueError(f"unsupported to_data_format {to_data_format}")
@staticmethod @staticmethod
def getCurrentDeviceConfig(): def getCurrentDeviceConfig():
if nn.current_DeviceConfig is None: if nn.current_DeviceConfig is None:
@ -159,6 +234,13 @@ class nn():
nn.tf_sess.close() nn.tf_sess.close()
nn.tf_sess = None nn.tf_sess = None
@staticmethod
def tf_get_current_device():
# Undocumented access to last tf.device(...)
objs = nn.tf.get_default_graph()._device_function_stack.peek_objs()
if len(objs) != 0:
return objs[0].display_name
return nn.tf_default_device
@staticmethod @staticmethod
def ask_choose_device_idxs(choose_only_one=False, allow_cpu=True, suggest_best_multi_gpu=False, suggest_all_gpu=False, return_device_config=False): def ask_choose_device_idxs(choose_only_one=False, allow_cpu=True, suggest_best_multi_gpu=False, suggest_all_gpu=False, return_device_config=False):

16
core/leras/optimizers.py
View file

@ -73,7 +73,7 @@ def initialize_optimizers(nn):
e = tf.device('/CPU:0') if vars_on_cpu else None e = tf.device('/CPU:0') if vars_on_cpu else None
if e: e.__enter__() if e: e.__enter__()
with tf.variable_scope(self.name): with tf.variable_scope(self.name):
accumulators = [ tf.get_variable ( f'acc_{i+self.accumulator_counter}', v.shape, initializer=tf.initializers.constant(0.0), trainable=False) accumulators = [ tf.get_variable ( f'acc_{i+self.accumulator_counter}', v.shape, dtype=v.dtype, initializer=tf.initializers.constant(0.0), trainable=False)
for (i, v ) in enumerate(trainable_weights) ] for (i, v ) in enumerate(trainable_weights) ]
self.accumulators_dict.update ( { v.name : acc for v,acc in zip(trainable_weights,accumulators) } ) self.accumulators_dict.update ( { v.name : acc for v,acc in zip(trainable_weights,accumulators) } )
@ -81,13 +81,13 @@ def initialize_optimizers(nn):
self.accumulator_counter += len(trainable_weights) self.accumulator_counter += len(trainable_weights)
if self.lr_dropout != 1.0: if self.lr_dropout != 1.0:
lr_rnds = [ nn.tf_random_binomial( v.shape, p=self.lr_dropout) for v in trainable_weights ] lr_rnds = [ nn.tf_random_binomial( v.shape, p=self.lr_dropout, dtype=v.dtype) for v in trainable_weights ]
self.lr_rnds_dict.update ( { v.name : rnd for v,rnd in zip(trainable_weights,lr_rnds) } ) self.lr_rnds_dict.update ( { v.name : rnd for v,rnd in zip(trainable_weights,lr_rnds) } )
if e: e.__exit__(None, None, None) if e: e.__exit__(None, None, None)
def get_update_op(self, grads_vars): def get_update_op(self, grads_vars):
updates = [] updates = []
lr = self.lr
if self.clipnorm > 0.0: if self.clipnorm > 0.0:
norm = tf.sqrt( sum([tf.reduce_sum(tf.square(g)) for g,v in grads_vars])) norm = tf.sqrt( sum([tf.reduce_sum(tf.square(g)) for g,v in grads_vars]))
updates += [ state_ops.assign_add( self.iterations, 1) ] updates += [ state_ops.assign_add( self.iterations, 1) ]
@ -96,8 +96,14 @@ def initialize_optimizers(nn):
g = self.tf_clip_norm(g, self.clipnorm, norm) g = self.tf_clip_norm(g, self.clipnorm, norm)
a = self.accumulators_dict[v.name] a = self.accumulators_dict[v.name]
new_a = self.rho * a + (1. - self.rho) * tf.square(g)
v_diff = - lr * g / (tf.sqrt(new_a) + self.epsilon) rho = tf.cast(self.rho, a.dtype)
new_a = rho * a + (1. - rho) * tf.square(g)
lr = tf.cast(self.lr, a.dtype)
epsilon = tf.cast(self.epsilon, a.dtype)
v_diff = - lr * g / (tf.sqrt(new_a) + epsilon)
if self.lr_dropout != 1.0: if self.lr_dropout != 1.0:
lr_rnd = self.lr_rnds_dict[v.name] lr_rnd = self.lr_rnds_dict[v.name]
v_diff *= lr_rnd v_diff *= lr_rnd

205
core/leras/tensor_ops.py
View file

@ -40,6 +40,9 @@ def initialize_tensor_ops(nn):
nn.tf_gradients = tf_gradients nn.tf_gradients = tf_gradients
def tf_average_gv_list(grad_var_list, tf_device_string=None): def tf_average_gv_list(grad_var_list, tf_device_string=None):
if len(grad_var_list) == 1:
return grad_var_list[0]
e = tf.device(tf_device_string) if tf_device_string is not None else None e = tf.device(tf_device_string) if tf_device_string is not None else None
if e is not None: e.__enter__() if e is not None: e.__enter__()
result = [] result = []
@ -58,6 +61,9 @@ def initialize_tensor_ops(nn):
nn.tf_average_gv_list = tf_average_gv_list nn.tf_average_gv_list = tf_average_gv_list
def tf_average_tensor_list(tensors_list, tf_device_string=None): def tf_average_tensor_list(tensors_list, tf_device_string=None):
if len(tensors_list) == 1:
return tensors_list[0]
e = tf.device(tf_device_string) if tf_device_string is not None else None e = tf.device(tf_device_string) if tf_device_string is not None else None
if e is not None: e.__enter__() if e is not None: e.__enter__()
result = tf.reduce_mean(tf.concat ([tf.expand_dims(t, 0) for t in tensors_list], 0), 0) result = tf.reduce_mean(tf.concat ([tf.expand_dims(t, 0) for t in tensors_list], 0), 0)
@ -65,36 +71,14 @@ def initialize_tensor_ops(nn):
return result return result
nn.tf_average_tensor_list = tf_average_tensor_list nn.tf_average_tensor_list = tf_average_tensor_list
def tf_dot(x, y): def tf_concat (tensors_list, axis):
if x.shape.ndims > 2 or y.shape.ndims > 2: """
x_shape = [] Better version.
for i, s in zip( x.shape.as_list(), array_ops.unstack(array_ops.shape(x))): """
if i is not None: if len(tensors_list) == 1:
x_shape.append(i) return tensors_list[0]
else: return tf.concat(tensors_list, axis)
x_shape.append(s) nn.tf_concat = tf_concat
x_shape = tuple(x_shape)
y_shape = []
for i, s in zip( y.shape.as_list(), array_ops.unstack(array_ops.shape(y))):
if i is not None:
y_shape.append(i)
else:
y_shape.append(s)
y_shape = tuple(y_shape)
y_permute_dim = list(range(y.shape.ndims))
y_permute_dim = [y_permute_dim.pop(-2)] + y_permute_dim
xt = array_ops.reshape(x, [-1, x_shape[-1]])
yt = array_ops.reshape(array_ops.transpose(y, perm=y_permute_dim), [y_shape[-2], -1])
import code
code.interact(local=dict(globals(), **locals()))
return array_ops.reshape(math_ops.matmul(xt, yt), x_shape[:-1] + y_shape[:-2] + y_shape[-1:])
if isinstance(x, sparse_tensor.SparseTensor):
out = sparse_ops.sparse_tensor_dense_matmul(x, y)
else:
out = math_ops.matmul(x, y)
return out
nn.tf_dot = tf_dot
def tf_gelu(x): def tf_gelu(x):
cdf = 0.5 * (1.0 + tf.nn.tanh((np.sqrt(2 / np.pi) * (x + 0.044715 * tf.pow(x, 3))))) cdf = 0.5 * (1.0 + tf.nn.tanh((np.sqrt(2 / np.pi) * (x + 0.044715 * tf.pow(x, 3)))))
@ -102,6 +86,13 @@ def initialize_tensor_ops(nn):
nn.tf_gelu = tf_gelu nn.tf_gelu = tf_gelu
def tf_upsample2d(x, size=2): def tf_upsample2d(x, size=2):
if nn.data_format == "NCHW":
b,c,h,w = x.shape.as_list()
x = tf.reshape (x, (-1,c,h,1,w,1) )
x = tf.tile(x, (1,1,1,size,1,size) )
x = tf.reshape (x, (-1,c,h*size,w*size) )
return x
else:
return tf.image.resize_nearest_neighbor(x, (x.shape[1]*size, x.shape[2]*size) ) return tf.image.resize_nearest_neighbor(x, (x.shape[1]*size, x.shape[2]*size) )
nn.tf_upsample2d = tf_upsample2d nn.tf_upsample2d = tf_upsample2d
@ -109,18 +100,24 @@ def initialize_tensor_ops(nn):
return tf.image.resize_images(x, (x.shape[1]*size, x.shape[2]*size) ) return tf.image.resize_images(x, (x.shape[1]*size, x.shape[2]*size) )
nn.tf_upsample2d_bilinear = tf_upsample2d_bilinear nn.tf_upsample2d_bilinear = tf_upsample2d_bilinear
def tf_flatten(x, dynamic_dims=False): def tf_flatten(x):
""" if nn.data_format == "NHWC":
dynamic_dims allows to flatten without knowing size on input dims # match NCHW version in order to switch data_format without problems
""" x = tf.transpose(x, (0,3,1,2) )
if dynamic_dims:
sh = tf.shape(x)
return tf.reshape (x, (sh[0], tf.reduce_prod(sh[1:]) ) )
else:
return tf.reshape (x, (-1, np.prod(x.shape[1:])) ) return tf.reshape (x, (-1, np.prod(x.shape[1:])) )
nn.tf_flatten = tf_flatten nn.tf_flatten = tf_flatten
def tf_reshape_4D(x, w,h,c):
if nn.data_format == "NHWC":
# match NCHW version in order to switch data_format without problems
x = tf.reshape (x, (-1,c,h,w))
x = tf.transpose(x, (0,2,3,1) )
return x
else:
return tf.reshape (x, (-1,c,h,w))
nn.tf_reshape_4D = tf_reshape_4D
def tf_random_binomial(shape, p=0.0, dtype=None, seed=None): def tf_random_binomial(shape, p=0.0, dtype=None, seed=None):
if dtype is None: if dtype is None:
dtype=tf.float32 dtype=tf.float32
@ -142,41 +139,42 @@ def initialize_tensor_ops(nn):
kernel_1d = np.array([gaussian(x, mean, sigma) for x in range(kernel_size)]) kernel_1d = np.array([gaussian(x, mean, sigma) for x in range(kernel_size)])
np_kernel = np.outer(kernel_1d, kernel_1d).astype(np.float32) np_kernel = np.outer(kernel_1d, kernel_1d).astype(np.float32)
kernel = np_kernel / np.sum(np_kernel) kernel = np_kernel / np.sum(np_kernel)
return kernel return kernel, kernel_size
gauss_kernel = make_kernel(radius) gauss_kernel, kernel_size = make_kernel(radius)
gauss_kernel = gauss_kernel[:, :,np.newaxis, np.newaxis] padding = kernel_size//2
kernel_size = gauss_kernel.shape[0] if padding != 0:
if nn.data_format == "NHWC":
inputs = [ input[:,:,:,i:i+1] for i in range( input.shape[-1] ) ] padding = [ [0,0], [padding,padding], [padding,padding], [0,0] ]
else:
padding = [ [0,0], [0,0], [padding,padding], [padding,padding] ]
else:
padding = None
gauss_kernel = gauss_kernel[:,:,None,None]
outputs = [] outputs = []
for i in range(len(inputs)): for i in range(input.shape[nn.conv2d_ch_axis]):
x = inputs[i] x = input[:,:,:,i:i+1] if nn.data_format == "NHWC" \
if kernel_size != 0: else input[:,i:i+1,:,:]
padding = kernel_size//2
x = tf.pad (x, [ [0,0], [padding,padding], [padding,padding], [0,0] ] )
outputs += [ tf.nn.conv2d(x, tf.constant(gauss_kernel, dtype=nn.tf_floatx ) , strides=[1,1,1,1], padding="VALID") ] if padding is not None:
x = tf.pad (x, padding)
outputs += [ tf.nn.conv2d(x, tf.constant(gauss_kernel, dtype=input.dtype ), strides=[1,1,1,1], padding="VALID", data_format=nn.data_format) ]
return tf.concat (outputs, axis=-1) return tf.concat (outputs, axis=nn.conv2d_ch_axis)
nn.tf_gaussian_blur = tf_gaussian_blur nn.tf_gaussian_blur = tf_gaussian_blur
def tf_style_loss(target, style, gaussian_blur_radius=0.0, loss_weight=1.0, step_size=1): def tf_style_loss(target, style, gaussian_blur_radius=0.0, loss_weight=1.0, step_size=1):
def sd(content, style, loss_weight): def sd(content, style, loss_weight):
content_nc = content.shape[-1] content_nc = content.shape[ nn.conv2d_ch_axis ]
style_nc = style.shape[-1] style_nc = style.shape[nn.conv2d_ch_axis]
if content_nc != style_nc: if content_nc != style_nc:
raise Exception("style_loss() content_nc != style_nc") raise Exception("style_loss() content_nc != style_nc")
c_mean, c_var = tf.nn.moments(content, axes=nn.conv2d_spatial_axes, keep_dims=True)
axes = [1,2] s_mean, s_var = tf.nn.moments(style, axes=nn.conv2d_spatial_axes, keep_dims=True)
c_mean, c_var = tf.nn.moments(content, axes=axes, keep_dims=True)
s_mean, s_var = tf.nn.moments(style, axes=axes, keep_dims=True)
c_std, s_std = tf.sqrt(c_var + 1e-5), tf.sqrt(s_var + 1e-5) c_std, s_std = tf.sqrt(c_var + 1e-5), tf.sqrt(s_var + 1e-5)
mean_loss = tf.reduce_sum(tf.square(c_mean-s_mean), axis=[1,2,3]) mean_loss = tf.reduce_sum(tf.square(c_mean-s_mean), axis=[1,2,3])
std_loss = tf.reduce_sum(tf.square(c_std-s_std), axis=[1,2,3]) std_loss = tf.reduce_sum(tf.square(c_std-s_std), axis=[1,2,3])
return (mean_loss + std_loss) * ( loss_weight / content_nc.value ) return (mean_loss + std_loss) * ( loss_weight / content_nc.value )
if gaussian_blur_radius > 0.0: if gaussian_blur_radius > 0.0:
@ -187,46 +185,29 @@ def initialize_tensor_ops(nn):
nn.tf_style_loss = tf_style_loss nn.tf_style_loss = tf_style_loss
def tf_channel_histogram (input, bins, data_range):
range_min, range_max = data_range
bin_range = (range_max-range_min) / (bins-1)
reduce_axes = [*range(input.shape.ndims)][1:]
x = input
x += bin_range/2
output = []
for i in range(bins-1, -1, -1):
y = x - (i*bin_range)
ones_mask = tf.sign( tf.nn.relu(y) )
x = x * (1.0 - ones_mask)
output.append ( tf.expand_dims(tf.reduce_sum (ones_mask, axis=reduce_axes ), -1) )
return tf.concat(output[::-1],-1)
nn.tf_channel_histogram = tf_channel_histogram
def tf_histogram(input, bins=256, data_range=(0,1.0)):
return tf.concat ( [tf.expand_dims( tf_channel_histogram( input[...,i], bins=bins, data_range=data_range ), -1 ) for i in range(input.shape[-1])], -1 )
nn.tf_histogram = tf_histogram
def tf_dssim(img1,img2, max_val, filter_size=11, filter_sigma=1.5, k1=0.01, k2=0.03): def tf_dssim(img1,img2, max_val, filter_size=11, filter_sigma=1.5, k1=0.01, k2=0.03):
if img1.dtype != img2.dtype:
raise ValueError("img1.dtype != img2.dtype")
ch = img2.shape[-1] not_float32 = img1.dtype != tf.float32
def _fspecial_gauss(size, sigma): if not_float32:
#Function to mimic the 'fspecial' gaussian MATLAB function. img_dtype = img1.dtype
coords = np.arange(0, size, dtype=nn.np_floatx) img1 = tf.cast(img1, tf.float32)
coords -= (size - 1 ) / 2.0 img2 = tf.cast(img2, tf.float32)
g = coords**2
g *= ( -0.5 / (sigma**2) )
g = np.reshape (g, (1,-1)) + np.reshape(g, (-1,1) )
g = tf.constant ( np.reshape (g, (1,-1)), dtype=nn.tf_floatx )
g = tf.nn.softmax(g)
g = tf.reshape (g, (size, size, 1, 1))
g = tf.tile (g, (1,1,ch,1))
return g
kernel = _fspecial_gauss(filter_size,filter_sigma) kernel = np.arange(0, filter_size, dtype=np.float32)
kernel -= (filter_size - 1 ) / 2.0
kernel = kernel**2
kernel *= ( -0.5 / (filter_sigma**2) )
kernel = np.reshape (kernel, (1,-1)) + np.reshape(kernel, (-1,1) )
kernel = tf.constant ( np.reshape (kernel, (1,-1)), dtype=tf.float32 )
kernel = tf.nn.softmax(kernel)
kernel = tf.reshape (kernel, (filter_size, filter_size, 1, 1))
kernel = tf.tile (kernel, (1,1, img1.shape[ nn.conv2d_ch_axis ] ,1))
def reducer(x): def reducer(x):
return tf.nn.depthwise_conv2d(x, kernel, strides=[1,1,1,1], padding='VALID') return tf.nn.depthwise_conv2d(x, kernel, strides=[1,1,1,1], padding='VALID', data_format=nn.data_format)
c1 = (k1 * max_val) ** 2 c1 = (k1 * max_val) ** 2
c2 = (k2 * max_val) ** 2 c2 = (k2 * max_val) ** 2
@ -242,10 +223,44 @@ def initialize_tensor_ops(nn):
c2 *= 1.0 #compensation factor c2 *= 1.0 #compensation factor
cs = (num1 - num0 + c2) / (den1 - den0 + c2) cs = (num1 - num0 + c2) / (den1 - den0 + c2)
ssim_val = tf.reduce_mean(luminance * cs, axis=(-3, -2) ) ssim_val = tf.reduce_mean(luminance * cs, axis=nn.conv2d_spatial_axes )
return(1.0 - ssim_val ) / 2.0 dssim = (1.0 - ssim_val ) / 2.0
if not_float32:
dssim = tf.cast(dssim, img_dtype)
return dssim
nn.tf_dssim = tf_dssim nn.tf_dssim = tf_dssim
def tf_space_to_depth(x, size):
if nn.data_format == "NHWC":
# match NCHW version in order to switch data_format without problems
b,h,w,c = x.shape.as_list()
oh, ow = h // size, w // size
x = tf.reshape(x, (-1, size, oh, size, ow, c))
x = tf.transpose(x, (0, 2, 4, 1, 3, 5))
x = tf.reshape(x, (-1, oh, ow, size* size* c ))
return x
else:
return tf.space_to_depth(x, size, data_format=nn.data_format)
nn.tf_space_to_depth = tf_space_to_depth
def tf_depth_to_space(x, size):
if nn.data_format == "NHWC":
# match NCHW version in order to switch data_format without problems
b,h,w,c = x.shape.as_list()
oh, ow = h * size, w * size
oc = c // (size * size)
x = tf.reshape(x, (-1, h, w, size, size, oc, ) )
x = tf.transpose(x, (0, 1, 3, 2, 4, 5))
x = tf.reshape(x, (-1, oh, ow, oc, ))
return x
else:
return tf.depth_to_space(x, size, data_format=nn.data_format)
nn.tf_depth_to_space = tf_depth_to_space
def tf_rgb_to_lab(srgb): def tf_rgb_to_lab(srgb):
srgb_pixels = tf.reshape(srgb, [-1, 3]) srgb_pixels = tf.reshape(srgb, [-1, 3])
linear_mask = tf.cast(srgb_pixels <= 0.04045, dtype=tf.float32) linear_mask = tf.cast(srgb_pixels <= 0.04045, dtype=tf.float32)

15
facelib/FANExtractor.py
View file

@ -18,7 +18,7 @@ class FANExtractor(object):
if not model_path.exists(): if not model_path.exists():
raise Exception("Unable to load FANExtractor model") raise Exception("Unable to load FANExtractor model")
nn.initialize() nn.initialize(data_format="NHWC")
tf = nn.tf tf = nn.tf
class ConvBlock(nn.ModelBase): class ConvBlock(nn.ModelBase):
@ -29,10 +29,10 @@ class FANExtractor(object):
self.bn1 = nn.BatchNorm2D(in_planes) self.bn1 = nn.BatchNorm2D(in_planes)
self.conv1 = nn.Conv2D (in_planes, out_planes/2, kernel_size=3, strides=1, padding='SAME', use_bias=False ) self.conv1 = nn.Conv2D (in_planes, out_planes/2, kernel_size=3, strides=1, padding='SAME', use_bias=False )
self.bn2 = nn.BatchNorm2D(out_planes/2) self.bn2 = nn.BatchNorm2D(out_planes//2)
self.conv2 = nn.Conv2D (out_planes/2, out_planes/4, kernel_size=3, strides=1, padding='SAME', use_bias=False ) self.conv2 = nn.Conv2D (out_planes/2, out_planes/4, kernel_size=3, strides=1, padding='SAME', use_bias=False )
self.bn3 = nn.BatchNorm2D(out_planes/4) self.bn3 = nn.BatchNorm2D(out_planes//4)
self.conv3 = nn.Conv2D (out_planes/4, out_planes/4, kernel_size=3, strides=1, padding='SAME', use_bias=False ) self.conv3 = nn.Conv2D (out_planes/4, out_planes/4, kernel_size=3, strides=1, padding='SAME', use_bias=False )
if self.in_planes != self.out_planes: if self.in_planes != self.out_planes:
@ -55,6 +55,7 @@ class FANExtractor(object):
x = self.bn3(x) x = self.bn3(x)
x = tf.nn.relu(x) x = tf.nn.relu(x)
x = out3 = self.conv3(x) x = out3 = self.conv3(x)
x = tf.concat ([out1, out2, out3], axis=-1) x = tf.concat ([out1, out2, out3], axis=-1)
if self.in_planes != self.out_planes: if self.in_planes != self.out_planes:
@ -148,7 +149,9 @@ class FANExtractor(object):
if i < 4 - 1: if i < 4 - 1:
ll = self.bl[i](ll) ll = self.bl[i](ll)
previous = previous + ll + self.al[i](tmp_out) previous = previous + ll + self.al[i](tmp_out)
return outputs[-1] x = outputs[-1]
x = tf.transpose(x, (0,3,1,2) )
return x
e = None e = None
if place_model_on_cpu: if place_model_on_cpu:
@ -159,7 +162,7 @@ class FANExtractor(object):
self.model.load_weights(str(model_path)) self.model.load_weights(str(model_path))
if e is not None: e.__exit__(None,None,None) if e is not None: e.__exit__(None,None,None)
self.model.build_for_run ([ ( tf.float32, (256,256,3) ) ]) self.model.build_for_run ([ ( tf.float32, (None,256,256,3) ) ])
def extract (self, input_image, rects, second_pass_extractor=None, is_bgr=True, multi_sample=False): def extract (self, input_image, rects, second_pass_extractor=None, is_bgr=True, multi_sample=False):
if len(rects) == 0: if len(rects) == 0:
@ -197,7 +200,7 @@ class FANExtractor(object):
predicted = [] predicted = []
for i in range( len(images) ): for i in range( len(images) ):
predicted += [ self.model.run ( [ images[i][None,...] ] ).transpose (0,3,1,2)[0] ] predicted += [ self.model.run ( [ images[i][None,...] ] )[0] ]
predicted = np.stack(predicted) predicted = np.stack(predicted)

8
facelib/FaceEnhancer.py
View file

@ -11,7 +11,7 @@ class FaceEnhancer(object):
x4 face enhancer x4 face enhancer
""" """
def __init__(self, place_model_on_cpu=False): def __init__(self, place_model_on_cpu=False):
nn.initialize() nn.initialize(data_format="NHWC")
tf = nn.tf tf = nn.tf
class FaceEnhancer (nn.ModelBase): class FaceEnhancer (nn.ModelBase):
@ -167,9 +167,9 @@ class FaceEnhancer(object):
self.model.load_weights (model_path) self.model.load_weights (model_path)
if e is not None: e.__exit__(None,None,None) if e is not None: e.__exit__(None,None,None)
self.model.build_for_run ([ (tf.float32, (192,192,3) ), self.model.build_for_run ([ (tf.float32, nn.get4Dshape (192,192,3) ),
(tf.float32, (1,) ), (tf.float32, (None,1,) ),
(tf.float32, (1,) ), (tf.float32, (None,1,) ),
]) ])

4
facelib/S3FDExtractor.py
View file

@ -8,7 +8,7 @@ from core.leras import nn
class S3FDExtractor(object): class S3FDExtractor(object):
def __init__(self, place_model_on_cpu=False): def __init__(self, place_model_on_cpu=False):
nn.initialize() nn.initialize(data_format="NHWC")
tf = nn.tf tf = nn.tf
model_path = Path(__file__).parent / "S3FD.npy" model_path = Path(__file__).parent / "S3FD.npy"
@ -167,7 +167,7 @@ class S3FDExtractor(object):
self.model.load_weights (model_path) self.model.load_weights (model_path)
if e is not None: e.__exit__(None,None,None) if e is not None: e.__exit__(None,None,None)
self.model.build_for_run ([ ( tf.float32, (None,None,3) ) ]) self.model.build_for_run ([ ( tf.float32, nn.get4Dshape (None,None,3) ) ])
def __enter__(self): def __enter__(self):
return self return self

125
facelib/TernausNet.py
View file

@ -20,7 +20,7 @@ TernausNet: U-Net with VGG11 Encoder Pre-Trained on ImageNet for Image Segmentat
class TernausNet(object): class TernausNet(object):
VERSION = 1 VERSION = 1
def __init__ (self, name, resolution, face_type_str, load_weights=True, weights_file_root=None, training=False, place_model_on_cpu=False): def __init__ (self, name, resolution, face_type_str, load_weights=True, weights_file_root=None, training=False, place_model_on_cpu=False):
nn.initialize() nn.initialize(data_format="NHWC")
tf = nn.tf tf = nn.tf
class Ternaus(nn.ModelBase): class Ternaus(nn.ModelBase):
@ -125,7 +125,7 @@ class TernausNet(object):
self.net.init_weights() self.net.init_weights()
if e is not None: e.__exit__(None,None,None) if e is not None: e.__exit__(None,None,None)
self.net.build_for_run ( [(tf.float32, (resolution,resolution,3))] ) self.net.build_for_run ( [(tf.float32, nn.get4Dshape (resolution,resolution,3) )] )
if training: if training:
raise Exception("training not supported yet") raise Exception("training not supported yet")
@ -195,124 +195,3 @@ class TernausNet(object):
result = result[0] result = result[0]
return result return result
"""
self.weights_path = weights_file_root / ('%s_%d_%s.h5' % (name, resolution, face_type_str) )
self.net.build()
self.net.features_0.set_weights ( self.model.get_layer('features.0').get_weights() )
self.net.features_3.set_weights ( self.model.get_layer('features.3').get_weights() )
self.net.features_6.set_weights ( self.model.get_layer('features.6').get_weights() )
self.net.features_8.set_weights ( self.model.get_layer('features.8').get_weights() )
self.net.features_11.set_weights ( self.model.get_layer('features.11').get_weights() )
self.net.features_13.set_weights ( self.model.get_layer('features.13').get_weights() )
self.net.features_16.set_weights ( self.model.get_layer('features.16').get_weights() )
self.net.features_18.set_weights ( self.model.get_layer('features.18').get_weights() )
self.net.conv_center.set_weights ( self.model.get_layer('CA.1').get_weights() )
self.net.conv1_up.set_weights ( self.model.get_layer('CA.2').get_weights() )
self.net.conv1.set_weights ( self.model.get_layer('CA.3').get_weights() )
self.net.conv2_up.set_weights ( self.model.get_layer('CA.4').get_weights() )
self.net.conv2.set_weights ( self.model.get_layer('CA.5').get_weights() )
self.net.conv3_up.set_weights ( self.model.get_layer('CA.6').get_weights() )
self.net.conv3.set_weights ( self.model.get_layer('CA.7').get_weights() )
self.net.conv4_up.set_weights ( self.model.get_layer('CA.8').get_weights() )
self.net.conv4.set_weights ( self.model.get_layer('CA.9').get_weights() )
self.net.conv5_up.set_weights ( self.model.get_layer('CA.10').get_weights() )
self.net.conv5.set_weights ( self.model.get_layer('CA.11').get_weights() )
self.net.out_conv.set_weights ( self.model.get_layer('CA.12').get_weights() )
self.net.build_for_run ( [ (tf.float32, (resolution,resolution,3)) ])
self.net.save_weights (self.weights_path2)
def extract (self, input_image):
input_shape_len = len(input_image.shape)
if input_shape_len == 3:
input_image = input_image[np.newaxis,...]
result = np.clip ( self.model.predict( [input_image] ), 0, 1.0 )
result[result < 0.1] = 0 #get rid of noise
if input_shape_len == 3:
result = result[0]
return result
@staticmethod
def BuildModel ( resolution, ngf=64):
exec( nn.initialize(), locals(), globals() )
inp = Input ( (resolution,resolution,3) )
x = inp
x = TernausNet.Flow(ngf=ngf)(x)
model = Model(inp,x)
return model
@staticmethod
def Flow(ngf=64):
exec( nn.initialize(), locals(), globals() )
def func(input):
x = input
x0 = x = Conv2D(ngf, kernel_size=3, strides=1, padding='same', activation='relu', name='features.0')(x)
x = BlurPool(filt_size=3)(x)
x1 = x = Conv2D(ngf*2, kernel_size=3, strides=1, padding='same', activation='relu', name='features.3')(x)
x = BlurPool(filt_size=3)(x)
x = Conv2D(ngf*4, kernel_size=3, strides=1, padding='same', activation='relu', name='features.6')(x)
x2 = x = Conv2D(ngf*4, kernel_size=3, strides=1, padding='same', activation='relu', name='features.8')(x)
x = BlurPool(filt_size=3)(x)
x = Conv2D(ngf*8, kernel_size=3, strides=1, padding='same', activation='relu', name='features.11')(x)
x3 = x = Conv2D(ngf*8, kernel_size=3, strides=1, padding='same', activation='relu', name='features.13')(x)
x = BlurPool(filt_size=3)(x)
x = Conv2D(ngf*8, kernel_size=3, strides=1, padding='same', activation='relu', name='features.16')(x)
x4 = x = Conv2D(ngf*8, kernel_size=3, strides=1, padding='same', activation='relu', name='features.18')(x)
x = BlurPool(filt_size=3)(x)
x = Conv2D(ngf*8, kernel_size=3, strides=1, padding='same', name='CA.1')(x)
x = Conv2DTranspose (ngf*4, 3, strides=2, padding='same', activation='relu', name='CA.2') (x)
x = Concatenate(axis=3)([ x, x4])
x = Conv2D (ngf*8, 3, strides=1, padding='same', activation='relu', name='CA.3') (x)
x = Conv2DTranspose (ngf*4, 3, strides=2, padding='same', activation='relu', name='CA.4') (x)
x = Concatenate(axis=3)([ x, x3])
x = Conv2D (ngf*8, 3, strides=1, padding='same', activation='relu', name='CA.5') (x)
x = Conv2DTranspose (ngf*2, 3, strides=2, padding='same', activation='relu', name='CA.6') (x)
x = Concatenate(axis=3)([ x, x2])
x = Conv2D (ngf*4, 3, strides=1, padding='same', activation='relu', name='CA.7') (x)
x = Conv2DTranspose (ngf, 3, strides=2, padding='same', activation='relu', name='CA.8') (x)
x = Concatenate(axis=3)([ x, x1])
x = Conv2D (ngf*2, 3, strides=1, padding='same', activation='relu', name='CA.9') (x)
x = Conv2DTranspose (ngf // 2, 3, strides=2, padding='same', activation='relu', name='CA.10') (x)
x = Concatenate(axis=3)([ x, x0])
x = Conv2D (ngf, 3, strides=1, padding='same', activation='relu', name='CA.11') (x)
return Conv2D(1, 3, strides=1, padding='same', activation='sigmoid', name='CA.12')(x)
return func
"""

2
mainscripts/Sorter.py
View file

@ -717,7 +717,7 @@ def sort_by_absdiff(input_path):
from core.leras import nn from core.leras import nn
device_config = nn.ask_choose_device_idxs(choose_only_one=True, return_device_config=True) device_config = nn.ask_choose_device_idxs(choose_only_one=True, return_device_config=True)
nn.initialize( device_config=device_config ) nn.initialize( device_config=device_config, data_format="NHWC" )
tf = nn.tf tf = nn.tf
image_paths = pathex.get_image_paths(input_path) image_paths = pathex.get_image_paths(input_path)

1
mainscripts/dev_misc.py
View file

@ -485,4 +485,3 @@ def dev_test(input_dir):
#import code #import code
#code.interact(local=dict(globals(), **locals())) #code.interact(local=dict(globals(), **locals()))

4
models/ModelBase.py
View file

@ -479,7 +479,7 @@ class ModelBase(object):
#Find the longest key name and value string. Used as column widths. #Find the longest key name and value string. Used as column widths.
width_name = max([len(k) for k in self.options.keys()] + [17]) + 1 # Single space buffer to left edge. Minimum of 17, the length of the longest static string used "Current iteration" width_name = max([len(k) for k in self.options.keys()] + [17]) + 1 # Single space buffer to left edge. Minimum of 17, the length of the longest static string used "Current iteration"
width_value = max([len(str(x)) for x in self.options.values()] + [len(str(self.get_iter())), len(self.get_model_name())]) + 1 # Single space buffer to right edge width_value = max([len(str(x)) for x in self.options.values()] + [len(str(self.get_iter())), len(self.get_model_name())]) + 1 # Single space buffer to right edge
if not self.device_config.cpu_only: #Check length of GPU names if len(self.device_config.devices) != 0: #Check length of GPU names
width_value = max([len(device.name)+1 for device in self.device_config.devices] + [width_value]) width_value = max([len(device.name)+1 for device in self.device_config.devices] + [width_value])
width_total = width_name + width_value + 2 #Plus 2 for ": " width_total = width_name + width_value + 2 #Plus 2 for ": "
@ -499,7 +499,7 @@ class ModelBase(object):
summary_text += [f'=={" Running On ":-^{width_total}}=='] # Training hardware info summary_text += [f'=={" Running On ":-^{width_total}}=='] # Training hardware info
summary_text += [f'=={" "*width_total}=='] summary_text += [f'=={" "*width_total}==']
if self.device_config.cpu_only: if len(self.device_config.devices) == 0:
summary_text += [f'=={"Using device": >{width_name}}: {"CPU": <{width_value}}=='] # cpu_only summary_text += [f'=={"Using device": >{width_name}}: {"CPU": <{width_value}}=='] # cpu_only
else: else:
for device in self.device_config.devices: for device in self.device_config.devices:

87
models/Model_Quick96/Model.py
View file

@ -13,10 +13,12 @@ from samplelib import *
class QModel(ModelBase): class QModel(ModelBase):
#override #override
def on_initialize(self): def on_initialize(self):
nn.initialize() device_config = nn.getCurrentDeviceConfig()
self.model_data_format = "NCHW" if len(device_config.devices) != 0 else "NHWC"
nn.initialize(data_format=self.model_data_format)
tf = nn.tf tf = nn.tf
conv_kernel_initializer = nn.initializers.ca conv_kernel_initializer = nn.initializers.ca()
class Downscale(nn.ModelBase): class Downscale(nn.ModelBase):
def __init__(self, in_ch, out_ch, kernel_size=5, dilations=1, subpixel=True, use_activator=True, *kwargs ): def __init__(self, in_ch, out_ch, kernel_size=5, dilations=1, subpixel=True, use_activator=True, *kwargs ):
@ -39,7 +41,7 @@ class QModel(ModelBase):
x = self.conv1(x) x = self.conv1(x)
if self.subpixel: if self.subpixel:
x = tf.nn.space_to_depth(x, 2) x = nn.tf_space_to_depth(x, 2)
if self.use_activator: if self.use_activator:
x = nn.tf_gelu(x) x = nn.tf_gelu(x)
@ -71,7 +73,7 @@ class QModel(ModelBase):
def forward(self, x): def forward(self, x):
x = self.conv1(x) x = self.conv1(x)
x = nn.tf_gelu(x) x = nn.tf_gelu(x)
x = tf.nn.depth_to_space(x, 2) x = nn.tf_depth_to_space(x, 2)
return x return x
class ResidualBlock(nn.ModelBase): class ResidualBlock(nn.ModelBase):
@ -109,7 +111,7 @@ class QModel(ModelBase):
def forward(self, inp): def forward(self, inp):
x = self.dense1(inp) x = self.dense1(inp)
x = self.dense2(x) x = self.dense2(x)
x = tf.reshape (x, (-1, lowest_dense_res, lowest_dense_res, self.ae_out_ch)) x = nn.tf_reshape_4D (x, lowest_dense_res, lowest_dense_res, self.ae_out_ch)
x = self.upscale1(x) x = self.upscale1(x)
x = self.res1(x) x = self.res1(x)
return x return x
@ -167,8 +169,8 @@ class QModel(ModelBase):
input_nc = 3 input_nc = 3
output_nc = 3 output_nc = 3
bgr_shape = (resolution, resolution, output_nc) bgr_shape = nn.get4Dshape(resolution,resolution,input_nc)
mask_shape = (resolution, resolution, 1) mask_shape = nn.get4Dshape(resolution,resolution,1)
lowest_dense_res = resolution // 16 lowest_dense_res = resolution // 16
self.model_filename_list = [] self.model_filename_list = []
@ -176,22 +178,22 @@ class QModel(ModelBase):
with tf.device ('/CPU:0'): with tf.device ('/CPU:0'):
#Place holders on CPU #Place holders on CPU
self.warped_src = tf.placeholder (tf.float32, (None,)+bgr_shape) self.warped_src = tf.placeholder (nn.tf_floatx, bgr_shape)
self.warped_dst = tf.placeholder (tf.float32, (None,)+bgr_shape) self.warped_dst = tf.placeholder (nn.tf_floatx, bgr_shape)
self.target_src = tf.placeholder (tf.float32, (None,)+bgr_shape) self.target_src = tf.placeholder (nn.tf_floatx, bgr_shape)
self.target_dst = tf.placeholder (tf.float32, (None,)+bgr_shape) self.target_dst = tf.placeholder (nn.tf_floatx, bgr_shape)
self.target_srcm = tf.placeholder (tf.float32, (None,)+mask_shape) self.target_srcm = tf.placeholder (nn.tf_floatx, mask_shape)
self.target_dstm = tf.placeholder (tf.float32, (None,)+mask_shape) self.target_dstm = tf.placeholder (nn.tf_floatx, mask_shape)
# Initializing model classes # Initializing model classes
with tf.device (models_opt_device): with tf.device (models_opt_device):
self.encoder = Encoder(in_ch=input_nc, e_ch=e_dims, name='encoder') self.encoder = Encoder(in_ch=input_nc, e_ch=e_dims, name='encoder')
encoder_out_ch = self.encoder.compute_output_shape ( (tf.float32, (None,resolution,resolution,input_nc)))[-1] 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, d_ch=d_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, d_ch=d_dims, name='inter')
inter_out_ch = self.inter.compute_output_shape ( (tf.float32, (None,encoder_out_ch)))[-1] 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, name='decoder_src') self.decoder_src = Decoder(in_ch=inter_out_ch, d_ch=d_dims, name='decoder_src')
self.decoder_dst = Decoder(in_ch=inter_out_ch, d_ch=d_dims, name='decoder_dst') self.decoder_dst = Decoder(in_ch=inter_out_ch, d_ch=d_dims, name='decoder_dst')
@ -271,11 +273,11 @@ class QModel(ModelBase):
gpu_src_loss = tf.reduce_mean ( 10*nn.tf_dssim(gpu_target_srcmasked_opt, gpu_pred_src_src_masked_opt, max_val=1.0, filter_size=int(resolution/11.6)), axis=[1]) gpu_src_loss = tf.reduce_mean ( 10*nn.tf_dssim(gpu_target_srcmasked_opt, gpu_pred_src_src_masked_opt, max_val=1.0, filter_size=int(resolution/11.6)), axis=[1])
gpu_src_loss += tf.reduce_mean ( 10*tf.square ( gpu_target_srcmasked_opt - gpu_pred_src_src_masked_opt ), axis=[1,2,3]) gpu_src_loss += tf.reduce_mean ( 10*tf.square ( gpu_target_srcmasked_opt - gpu_pred_src_src_masked_opt ), axis=[1,2,3])
gpu_src_loss += tf.reduce_mean ( tf.square( gpu_target_srcm - gpu_pred_src_srcm ),axis=[1,2,3] ) gpu_src_loss += tf.reduce_mean ( 10*tf.square( gpu_target_srcm - gpu_pred_src_srcm ),axis=[1,2,3] )
gpu_dst_loss = tf.reduce_mean ( 10*nn.tf_dssim(gpu_target_dst_masked_opt, gpu_pred_dst_dst_masked_opt, max_val=1.0, filter_size=int(resolution/11.6) ), axis=[1]) gpu_dst_loss = tf.reduce_mean ( 10*nn.tf_dssim(gpu_target_dst_masked_opt, gpu_pred_dst_dst_masked_opt, max_val=1.0, filter_size=int(resolution/11.6) ), axis=[1])
gpu_dst_loss += tf.reduce_mean ( 10*tf.square( gpu_target_dst_masked_opt- gpu_pred_dst_dst_masked_opt ), axis=[1,2,3]) gpu_dst_loss += tf.reduce_mean ( 10*tf.square( gpu_target_dst_masked_opt- gpu_pred_dst_dst_masked_opt ), axis=[1,2,3])
gpu_dst_loss += tf.reduce_mean ( tf.square( gpu_target_dstm - gpu_pred_dst_dstm ),axis=[1,2,3] ) gpu_dst_loss += tf.reduce_mean ( 10*tf.square( gpu_target_dstm - gpu_pred_dst_dstm ),axis=[1,2,3] )
gpu_src_losses += [gpu_src_loss] gpu_src_losses += [gpu_src_loss]
gpu_dst_losses += [gpu_dst_loss] gpu_dst_losses += [gpu_dst_loss]
@ -286,29 +288,16 @@ class QModel(ModelBase):
# Average losses and gradients, and create optimizer update ops # Average losses and gradients, and create optimizer update ops
with tf.device (models_opt_device): with tf.device (models_opt_device):
if gpu_count == 1: pred_src_src = nn.tf_concat(gpu_pred_src_src_list, 0)
pred_src_src = gpu_pred_src_src_list[0] pred_dst_dst = nn.tf_concat(gpu_pred_dst_dst_list, 0)
pred_dst_dst = gpu_pred_dst_dst_list[0] pred_src_dst = nn.tf_concat(gpu_pred_src_dst_list, 0)
pred_src_dst = gpu_pred_src_dst_list[0] pred_src_srcm = nn.tf_concat(gpu_pred_src_srcm_list, 0)
pred_src_srcm = gpu_pred_src_srcm_list[0] pred_dst_dstm = nn.tf_concat(gpu_pred_dst_dstm_list, 0)
pred_dst_dstm = gpu_pred_dst_dstm_list[0] pred_src_dstm = nn.tf_concat(gpu_pred_src_dstm_list, 0)
pred_src_dstm = gpu_pred_src_dstm_list[0]
src_loss = gpu_src_losses[0]
dst_loss = gpu_dst_losses[0]
src_dst_loss_gv = gpu_src_dst_loss_gvs[0]
else:
pred_src_src = tf.concat(gpu_pred_src_src_list, 0)
pred_dst_dst = tf.concat(gpu_pred_dst_dst_list, 0)
pred_src_dst = tf.concat(gpu_pred_src_dst_list, 0)
pred_src_srcm = tf.concat(gpu_pred_src_srcm_list, 0)
pred_dst_dstm = tf.concat(gpu_pred_dst_dstm_list, 0)
pred_src_dstm = tf.concat(gpu_pred_src_dstm_list, 0)
src_loss = nn.tf_average_tensor_list(gpu_src_losses) src_loss = nn.tf_average_tensor_list(gpu_src_losses)
dst_loss = nn.tf_average_tensor_list(gpu_dst_losses) dst_loss = nn.tf_average_tensor_list(gpu_dst_losses)
src_dst_loss_gv = nn.tf_average_gv_list (gpu_src_dst_loss_gvs) src_dst_loss_gv = nn.tf_average_gv_list (gpu_src_dst_loss_gvs)
src_dst_loss_gv_op = self.src_dst_opt.get_update_op (src_dst_loss_gv) src_dst_loss_gv_op = self.src_dst_opt.get_update_op (src_dst_loss_gv)
# Initializing training and view functions # Initializing training and view functions
@ -341,13 +330,11 @@ class QModel(ModelBase):
_, gpu_pred_dst_dstm = self.decoder_dst(gpu_dst_code) _, gpu_pred_dst_dstm = self.decoder_dst(gpu_dst_code)
def AE_merge( warped_dst): def AE_merge( warped_dst):
return nn.tf_sess.run ( [gpu_pred_src_dst, gpu_pred_dst_dstm, gpu_pred_src_dstm], feed_dict={self.warped_dst:warped_dst}) return nn.tf_sess.run ( [gpu_pred_src_dst, gpu_pred_dst_dstm, gpu_pred_src_dstm], feed_dict={self.warped_dst:warped_dst})
self.AE_merge = AE_merge self.AE_merge = AE_merge
# Loading/initializing all models/optimizers weights # Loading/initializing all models/optimizers weights
for model, filename in io.progress_bar_generator(self.model_filename_list, "Initializing models"): for model, filename in io.progress_bar_generator(self.model_filename_list, "Initializing models"):
do_init = self.is_first_run() do_init = self.is_first_run()
@ -368,7 +355,6 @@ class QModel(ModelBase):
model.init_weights() model.init_weights()
# initializing sample generators # initializing sample generators
if self.is_training: if self.is_training:
t = SampleProcessor.Types t = SampleProcessor.Types
face_type = t.FACE_TYPE_FULL face_type = t.FACE_TYPE_FULL
@ -384,16 +370,16 @@ class QModel(ModelBase):
self.set_training_data_generators ([ self.set_training_data_generators ([
SampleGeneratorFace(training_data_src_path, debug=self.is_debug(), batch_size=self.get_batch_size(), SampleGeneratorFace(training_data_src_path, debug=self.is_debug(), batch_size=self.get_batch_size(),
sample_process_options=SampleProcessor.Options(random_flip=True if self.pretrain else False), sample_process_options=SampleProcessor.Options(random_flip=True if self.pretrain else False),
output_sample_types = [ {'types' : (t.IMG_WARPED_TRANSFORMED, face_type, t.MODE_BGR), 'resolution':resolution, }, output_sample_types = [ {'types' : (t.IMG_WARPED_TRANSFORMED, face_type, t.MODE_BGR), 'data_format':nn.data_format, 'resolution':resolution, },
{'types' : (t.IMG_TRANSFORMED, face_type, t.MODE_BGR), 'resolution': resolution, }, {'types' : (t.IMG_TRANSFORMED, face_type, t.MODE_BGR), 'data_format':nn.data_format, 'resolution': resolution, },
{'types' : (t.IMG_TRANSFORMED, face_type, t.MODE_M), 'resolution': resolution } ], {'types' : (t.IMG_TRANSFORMED, face_type, t.MODE_M), 'data_format':nn.data_format, 'resolution': resolution } ],
generators_count=src_generators_count ), generators_count=src_generators_count ),
SampleGeneratorFace(training_data_dst_path, debug=self.is_debug(), batch_size=self.get_batch_size(), SampleGeneratorFace(training_data_dst_path, debug=self.is_debug(), batch_size=self.get_batch_size(),
sample_process_options=SampleProcessor.Options(random_flip=True if self.pretrain else False), sample_process_options=SampleProcessor.Options(random_flip=True if self.pretrain else False),
output_sample_types = [ {'types' : (t.IMG_WARPED_TRANSFORMED, face_type, t.MODE_BGR), 'resolution':resolution}, output_sample_types = [ {'types' : (t.IMG_WARPED_TRANSFORMED, face_type, t.MODE_BGR), 'data_format':nn.data_format, 'resolution':resolution},
{'types' : (t.IMG_TRANSFORMED, face_type, t.MODE_BGR), 'resolution': resolution}, {'types' : (t.IMG_TRANSFORMED, face_type, t.MODE_BGR), 'data_format':nn.data_format, 'resolution': resolution},
{'types' : (t.IMG_TRANSFORMED, face_type, t.MODE_M), 'resolution': resolution} ], {'types' : (t.IMG_TRANSFORMED, face_type, t.MODE_M), 'data_format':nn.data_format, 'resolution': resolution} ],
generators_count=dst_generators_count ) generators_count=dst_generators_count )
]) ])
@ -408,7 +394,6 @@ class QModel(ModelBase):
for model, filename in io.progress_bar_generator(self.get_model_filename_list(), "Saving", leave=False): for model, filename in io.progress_bar_generator(self.get_model_filename_list(), "Saving", leave=False):
model.save_weights ( self.get_strpath_storage_for_file(filename) ) model.save_weights ( self.get_strpath_storage_for_file(filename) )
#override #override
def onTrainOneIter(self): def onTrainOneIter(self):
if self.get_iter() % 3 == 0 and self.last_samples is not None: if self.get_iter() % 3 == 0 and self.last_samples is not None:
@ -435,9 +420,11 @@ class QModel(ModelBase):
[ [sample[0:n_samples] for sample in sample_list ] [ [sample[0:n_samples] for sample in sample_list ]
for sample_list in samples ] for sample_list in samples ]
S, D, SS, DD, DDM, SD, SDM = [ np.clip(x, 0.0, 1.0) for x in ([target_src,target_dst] + self.AE_view (target_src, target_dst) ) ] S, D, SS, DD, DDM, SD, SDM = [ np.clip( nn.to_data_format(x,"NHWC", self.model_data_format), 0.0, 1.0) for x in ([target_src,target_dst] + self.AE_view (target_src, target_dst) ) ]
DDM, SDM, = [ np.repeat (x, (3,), -1) for x in [DDM, SDM] ] DDM, SDM, = [ np.repeat (x, (3,), -1) for x in [DDM, SDM] ]
target_srcm, target_dstm = [ nn.to_data_format(x,"NHWC", self.model_data_format) for x in ([target_srcm, target_dstm] )]
result = [] result = []
st = [] st = []
for i in range(n_samples): for i in range(n_samples):
@ -456,8 +443,10 @@ class QModel(ModelBase):
return result return result
def predictor_func (self, face=None): def predictor_func (self, face=None):
face = face[None,...]
face = nn.to_data_format(face, self.model_data_format, "NHWC")
bgr, mask_dst_dstm, mask_src_dstm = self.AE_merge (face[np.newaxis,...]) bgr, mask_dst_dstm, mask_src_dstm = [ nn.to_data_format(x, "NHWC", self.model_data_format).astype(np.float32) for x in self.AE_merge (face) ]
mask = mask_dst_dstm[0] * mask_src_dstm[0] mask = mask_dst_dstm[0] * mask_src_dstm[0]
return bgr[0], mask[...,0] return bgr[0], mask[...,0]

166
models/Model_SAEHD/Model.py
View file

@ -26,14 +26,6 @@ class SAEHDModel(ModelBase):
suggest_batch_size = 4 suggest_batch_size = 4
yn_str = {True:'y',False:'n'} yn_str = {True:'y',False:'n'}
ask_override = self.ask_override()
if self.is_first_run() or ask_override:
self.ask_enable_autobackup()
self.ask_write_preview_history()
self.ask_target_iter()
self.ask_random_flip()
self.ask_batch_size(suggest_batch_size)
default_resolution = self.options['resolution'] = self.load_or_def_option('resolution', 128) default_resolution = self.options['resolution'] = self.load_or_def_option('resolution', 128)
default_face_type = self.options['face_type'] = self.load_or_def_option('face_type', 'f') default_face_type = self.options['face_type'] = self.load_or_def_option('face_type', 'f')
@ -47,26 +39,31 @@ class SAEHDModel(ModelBase):
default_d_mask_dims += default_d_mask_dims % 2 default_d_mask_dims += default_d_mask_dims % 2
default_d_mask_dims = self.options['d_mask_dims'] = self.load_or_def_option('d_mask_dims', default_d_mask_dims) default_d_mask_dims = self.options['d_mask_dims'] = self.load_or_def_option('d_mask_dims', default_d_mask_dims)
default_use_float16 = self.options['use_float16'] = self.load_or_def_option('use_float16', False)
default_learn_mask = self.options['learn_mask'] = self.load_or_def_option('learn_mask', True) default_learn_mask = self.options['learn_mask'] = self.load_or_def_option('learn_mask', True)
default_lr_dropout = self.options['lr_dropout'] = self.load_or_def_option('lr_dropout', False) default_lr_dropout = self.options['lr_dropout'] = self.load_or_def_option('lr_dropout', False)
default_random_warp = self.options['random_warp'] = self.load_or_def_option('random_warp', True) default_random_warp = self.options['random_warp'] = self.load_or_def_option('random_warp', True)
default_true_face_training = self.options['true_face_training'] = self.load_or_def_option('true_face_training', False) default_true_face_power = self.options['true_face_power'] = self.load_or_def_option('true_face_power', 0.0)
default_face_style_power = self.options['face_style_power'] = self.load_or_def_option('face_style_power', 0.0) default_face_style_power = self.options['face_style_power'] = self.load_or_def_option('face_style_power', 0.0)
default_bg_style_power = self.options['bg_style_power'] = self.load_or_def_option('bg_style_power', 0.0) default_bg_style_power = self.options['bg_style_power'] = self.load_or_def_option('bg_style_power', 0.0)
default_ct_mode = self.options['ct_mode'] = self.load_or_def_option('ct_mode', 'none') default_ct_mode = self.options['ct_mode'] = self.load_or_def_option('ct_mode', 'none')
default_clipgrad = self.options['clipgrad'] = self.load_or_def_option('clipgrad', False) default_clipgrad = self.options['clipgrad'] = self.load_or_def_option('clipgrad', False)
default_pretrain = self.options['pretrain'] = self.load_or_def_option('pretrain', False) default_pretrain = self.options['pretrain'] = self.load_or_def_option('pretrain', False)
ask_override = self.ask_override()
if self.is_first_run() or ask_override:
self.ask_enable_autobackup()
self.ask_write_preview_history()
self.ask_target_iter()
self.ask_random_flip()
self.ask_batch_size(suggest_batch_size)
if self.is_first_run(): if self.is_first_run():
resolution = io.input_int("Resolution", default_resolution, add_info="64-256", help_message="More resolution requires more VRAM and time to train. Value will be adjusted to multiple of 16.") resolution = io.input_int("Resolution", default_resolution, add_info="64-256", help_message="More resolution requires more VRAM and time to train. Value will be adjusted to multiple of 16.")
resolution = np.clip ( (resolution // 16) * 16, 64, 256) resolution = np.clip ( (resolution // 16) * 16, 64, 256)
self.options['resolution'] = resolution self.options['resolution'] = resolution
self.options['face_type'] = io.input_str ("Face type", default_face_type, ['h','mf','f'], help_message="Half / mid face / full face. Half face has better resolution, but covers less area of cheeks. Mid face is 30% wider than half face.").lower() self.options['face_type'] = io.input_str ("Face type", default_face_type, ['h','mf','f'], help_message="Half / mid face / full face. Half face has better resolution, but covers less area of cheeks. Mid face is 30% wider than half face.").lower()
if (self.is_first_run() or ask_override) and len(device_config.devices) == 1:
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.")
if self.is_first_run():
self.options['archi'] = io.input_str ("AE architecture", default_archi, ['dfhd','liaehd','df','liae'], help_message="'df' keeps faces more natural. 'liae' can fix overly different face shapes. 'hd' is heavyweight version for the best quality.").lower() #-s version is slower, but has decreased change to collapse. self.options['archi'] = io.input_str ("AE architecture", default_archi, ['dfhd','liaehd','df','liae'], help_message="'df' keeps faces more natural. 'liae' can fix overly different face shapes. 'hd' is heavyweight version for the best quality.").lower() #-s version is slower, but has decreased change to collapse.
self.options['ae_dims'] = np.clip ( io.input_int("AutoEncoder dimensions", default_ae_dims, add_info="32-1024", help_message="All face information will packed to AE dims. If amount of AE dims are not enough, then for example closed eyes will not be recognized. More dims are better, but require more VRAM. You can fine-tune model size to fit your GPU." ), 32, 1024 ) self.options['ae_dims'] = np.clip ( io.input_int("AutoEncoder dimensions", default_ae_dims, add_info="32-1024", help_message="All face information will packed to AE dims. If amount of AE dims are not enough, then for example closed eyes will not be recognized. More dims are better, but require more VRAM. You can fine-tune model size to fit your GPU." ), 32, 1024 )
@ -81,13 +78,19 @@ class SAEHDModel(ModelBase):
if self.is_first_run() or ask_override: if self.is_first_run() or ask_override:
self.options['learn_mask'] = io.input_bool ("Learn mask", default_learn_mask, help_message="Learning mask can help model to recognize face directions. Learn without mask can reduce model size, in this case merger forced to use 'not predicted mask' that is not smooth as predicted.") self.options['learn_mask'] = io.input_bool ("Learn mask", default_learn_mask, help_message="Learning mask can help model to recognize face directions. Learn without mask can reduce model size, in this case merger forced to use 'not predicted mask' that is not smooth as predicted.")
if self.is_first_run() or ask_override:
if len(device_config.devices) == 1:
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['use_float16'] = io.input_bool ("Use float16", default_use_float16, help_message="Experimental option. Reduces the model size by half. Increases the speed of training. Decreases the accuracy of the model. The model may collapse. Model does not study the mask in large resolutions.")
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 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 for less amount of iterations.")
if 'df' in self.options['archi']: if 'df' in self.options['archi']:
self.options['true_face_training'] = io.input_bool ("Enable 'true face' training", default_true_face_training, help_message="The result face will be more like src and will get extra sharpness. Enable it for last 10-20k iterations before conversion.") self.options['true_face_power'] = np.clip ( io.input_number (" 'True face' power.", default_true_face_power, add_info="0.0000 .. 1.0", help_message="Experimental option. Discriminates result face to be more like src face. Higher value - stronger discrimination. Comparison - https://i.imgur.com/czScS9q.png"), 0.0, 1.0 )
else: else:
self.options['true_face_training'] = False self.options['true_face_power'] = 0.0
self.options['face_style_power'] = np.clip ( io.input_number("Face style power", default_face_style_power, add_info="0.0..100.0", help_message="Learn to transfer face style details such as light and color conditions. Warning: Enable it only after 10k iters, when predicted face is clear enough to start learn style. Start from 0.1 value and check history changes. Enabling this option increases the chance of model collapse."), 0.0, 100.0 ) self.options['face_style_power'] = np.clip ( io.input_number("Face style power", default_face_style_power, add_info="0.0..100.0", help_message="Learn to transfer face style details such as light and color conditions. Warning: Enable it only after 10k iters, when predicted face is clear enough to start learn style. Start from 0.1 value and check history changes. Enabling this option increases the chance of model collapse."), 0.0, 100.0 )
self.options['bg_style_power'] = np.clip ( io.input_number("Background style power", default_bg_style_power, add_info="0.0..100.0", help_message="Learn to transfer background around face. This can make face more like dst. Enabling this option increases the chance of model collapse."), 0.0, 100.0 ) self.options['bg_style_power'] = np.clip ( io.input_number("Background style power", default_bg_style_power, add_info="0.0..100.0", help_message="Learn to transfer background around face. This can make face more like dst. Enabling this option increases the chance of model collapse."), 0.0, 100.0 )
@ -103,12 +106,16 @@ class SAEHDModel(ModelBase):
if self.pretrain_just_disabled: if self.pretrain_just_disabled:
self.set_iter(1) self.set_iter(1)
#override #override
def on_initialize(self): def on_initialize(self):
nn.initialize() device_config = nn.getCurrentDeviceConfig()
self.model_data_format = "NCHW" if len(device_config.devices) != 0 else "NHWC"
nn.initialize(floatx="float16" if self.options['use_float16'] else "float32",
data_format=self.model_data_format)
tf = nn.tf tf = nn.tf
conv_kernel_initializer = nn.initializers.ca conv_kernel_initializer = nn.initializers.ca()
class Downscale(nn.ModelBase): class Downscale(nn.ModelBase):
def __init__(self, in_ch, out_ch, kernel_size=5, dilations=1, subpixel=True, use_activator=True, *kwargs ): def __init__(self, in_ch, out_ch, kernel_size=5, dilations=1, subpixel=True, use_activator=True, *kwargs ):
@ -125,13 +132,13 @@ class SAEHDModel(ModelBase):
self.out_ch // (4 if self.subpixel else 1), self.out_ch // (4 if self.subpixel else 1),
kernel_size=self.kernel_size, kernel_size=self.kernel_size,
strides=1 if self.subpixel else 2, strides=1 if self.subpixel else 2,
padding='SAME', dilations=self.dilations, kernel_initializer=conv_kernel_initializer ) padding='SAME', dilations=self.dilations, kernel_initializer=conv_kernel_initializer)
def forward(self, x): def forward(self, x):
x = self.conv1(x) x = self.conv1(x)
if self.subpixel: if self.subpixel:
x = tf.nn.space_to_depth(x, 2) x = nn.tf_space_to_depth(x, 2)
if self.use_activator: if self.use_activator:
x = tf.nn.leaky_relu(x, 0.1) x = tf.nn.leaky_relu(x, 0.1)
@ -163,7 +170,7 @@ class SAEHDModel(ModelBase):
def forward(self, x): def forward(self, x):
x = self.conv1(x) x = self.conv1(x)
x = tf.nn.leaky_relu(x, 0.1) x = tf.nn.leaky_relu(x, 0.1)
x = tf.nn.depth_to_space(x, 2) x = nn.tf_depth_to_space(x, 2)
return x return x
class ResidualBlock(nn.ModelBase): class ResidualBlock(nn.ModelBase):
@ -211,7 +218,6 @@ class SAEHDModel(ModelBase):
nn.tf_flatten(self.down4(inp)) ], -1 ) nn.tf_flatten(self.down4(inp)) ], -1 )
else: else:
x = nn.tf_flatten(self.down1(inp)) x = nn.tf_flatten(self.down1(inp))
return x return x
class Inter(nn.ModelBase): class Inter(nn.ModelBase):
@ -222,14 +228,14 @@ class SAEHDModel(ModelBase):
def on_build(self): 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 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, kernel_initializer=tf.initializers.orthogonal ) self.dense1 = nn.Dense( in_ch, ae_ch )
self.dense2 = nn.Dense( ae_ch, lowest_dense_res * lowest_dense_res * ae_out_ch, kernel_initializer=tf.initializers.orthogonal ) self.dense2 = nn.Dense( ae_ch, lowest_dense_res * lowest_dense_res * ae_out_ch )
self.upscale1 = Upscale(ae_out_ch, ae_out_ch) self.upscale1 = Upscale(ae_out_ch, ae_out_ch)
def forward(self, inp): def forward(self, inp):
x = self.dense1(inp) x = self.dense1(inp)
x = self.dense2(x) x = self.dense2(x)
x = tf.reshape (x, (-1, lowest_dense_res, lowest_dense_res, self.ae_out_ch)) x = nn.tf_reshape_4D (x, lowest_dense_res, lowest_dense_res, self.ae_out_ch)
x = self.upscale1(x) x = self.upscale1(x)
return x return x
@ -307,7 +313,7 @@ class SAEHDModel(ModelBase):
class CodeDiscriminator(nn.ModelBase): class CodeDiscriminator(nn.ModelBase):
def on_build(self, in_ch, code_res, ch=256): def on_build(self, in_ch, code_res, ch=256):
n_downscales = 2 + code_res // 8 n_downscales = 1 + code_res // 8
self.convs = [] self.convs = []
prev_ch = in_ch prev_ch = in_ch
@ -343,8 +349,8 @@ class SAEHDModel(ModelBase):
input_nc = 3 input_nc = 3
output_nc = 3 output_nc = 3
bgr_shape = (resolution, resolution, output_nc) bgr_shape = nn.get4Dshape(resolution,resolution,input_nc)
mask_shape = (resolution, resolution, 1) mask_shape = nn.get4Dshape(resolution,resolution,1)
lowest_dense_res = resolution // 16 lowest_dense_res = resolution // 16
self.model_filename_list = [] self.model_filename_list = []
@ -352,23 +358,23 @@ class SAEHDModel(ModelBase):
with tf.device ('/CPU:0'): with tf.device ('/CPU:0'):
#Place holders on CPU #Place holders on CPU
self.warped_src = tf.placeholder (tf.float32, (None,)+bgr_shape) self.warped_src = tf.placeholder (nn.tf_floatx, bgr_shape)
self.warped_dst = tf.placeholder (tf.float32, (None,)+bgr_shape) self.warped_dst = tf.placeholder (nn.tf_floatx, bgr_shape)
self.target_src = tf.placeholder (tf.float32, (None,)+bgr_shape) self.target_src = tf.placeholder (nn.tf_floatx, bgr_shape)
self.target_dst = tf.placeholder (tf.float32, (None,)+bgr_shape) self.target_dst = tf.placeholder (nn.tf_floatx, bgr_shape)
self.target_srcm = tf.placeholder (tf.float32, (None,)+mask_shape) self.target_srcm = tf.placeholder (nn.tf_floatx, mask_shape)
self.target_dstm = tf.placeholder (tf.float32, (None,)+mask_shape) self.target_dstm = tf.placeholder (nn.tf_floatx, mask_shape)
# Initializing model classes # Initializing model classes
with tf.device (models_opt_device): with tf.device (models_opt_device):
if 'df' in archi: if 'df' in archi:
self.encoder = Encoder(in_ch=input_nc, e_ch=e_dims, is_hd='hd' in archi, name='encoder') self.encoder = Encoder(in_ch=input_nc, e_ch=e_dims, is_hd='hd' in archi, name='encoder')
encoder_out_ch = self.encoder.compute_output_shape ( (tf.float32, (None,resolution,resolution,input_nc)))[-1] 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, name='inter')
inter_out_ch = self.inter.compute_output_shape ( (tf.float32, (None,encoder_out_ch)))[-1] 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='hd' in archi, 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='hd' in archi, name='decoder_dst')
@ -379,21 +385,20 @@ class SAEHDModel(ModelBase):
[self.decoder_dst, 'decoder_dst.npy'] ] [self.decoder_dst, 'decoder_dst.npy'] ]
if self.is_training: if self.is_training:
if self.options['true_face_training']: if self.options['true_face_power'] != 0:
self.dis = CodeDiscriminator(ae_dims, code_res=lowest_dense_res*2, name='dis' ) self.dis = CodeDiscriminator(ae_dims, code_res=lowest_dense_res*2, name='dis' )
self.model_filename_list += [ [self.dis, 'dis.npy'] ] self.model_filename_list += [ [self.dis, 'dis.npy'] ]
elif 'liae' in archi: elif 'liae' in archi:
self.encoder = Encoder(in_ch=input_nc, e_ch=e_dims, is_hd='hd' in archi, name='encoder') self.encoder = Encoder(in_ch=input_nc, e_ch=e_dims, is_hd='hd' in archi, name='encoder')
encoder_out_ch = self.encoder.compute_output_shape ( (tf.float32, (None,resolution,resolution,input_nc)))[-1] 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, 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, name='inter_B')
inter_AB_out_ch = self.inter_AB.compute_output_shape ( (tf.float32, (None,encoder_out_ch)))[-1] 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_shape ( (tf.float32, (None,encoder_out_ch)))[-1] 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 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='hd' in archi, name='decoder')
self.model_filename_list += [ [self.encoder, 'encoder.npy'], self.model_filename_list += [ [self.encoder, 'encoder.npy'],
@ -418,7 +423,7 @@ class SAEHDModel(ModelBase):
self.src_dst_opt.initialize_variables (self.src_dst_all_trainable_weights, vars_on_cpu=optimizer_vars_on_cpu) self.src_dst_opt.initialize_variables (self.src_dst_all_trainable_weights, vars_on_cpu=optimizer_vars_on_cpu)
if self.options['true_face_training']: if self.options['true_face_power'] != 0:
self.D_opt = nn.TFRMSpropOptimizer(lr=lr, lr_dropout=lr_dropout, clipnorm=clipnorm, name='D_opt') self.D_opt = nn.TFRMSpropOptimizer(lr=lr, lr_dropout=lr_dropout, clipnorm=clipnorm, name='D_opt')
self.D_opt.initialize_variables ( self.dis.get_weights(), vars_on_cpu=optimizer_vars_on_cpu) self.D_opt.initialize_variables ( self.dis.get_weights(), vars_on_cpu=optimizer_vars_on_cpu)
self.model_filename_list += [ (self.D_opt, 'D_opt.npy') ] self.model_filename_list += [ (self.D_opt, 'D_opt.npy') ]
@ -466,12 +471,12 @@ class SAEHDModel(ModelBase):
elif 'liae' in archi: elif 'liae' in archi:
gpu_src_code = self.encoder (gpu_warped_src) gpu_src_code = self.encoder (gpu_warped_src)
gpu_src_inter_AB_code = self.inter_AB (gpu_src_code) gpu_src_inter_AB_code = self.inter_AB (gpu_src_code)
gpu_src_code = tf.concat([gpu_src_inter_AB_code,gpu_src_inter_AB_code],-1) gpu_src_code = tf.concat([gpu_src_inter_AB_code,gpu_src_inter_AB_code], nn.conv2d_ch_axis )
gpu_dst_code = self.encoder (gpu_warped_dst) gpu_dst_code = self.encoder (gpu_warped_dst)
gpu_dst_inter_B_code = self.inter_B (gpu_dst_code) gpu_dst_inter_B_code = self.inter_B (gpu_dst_code)
gpu_dst_inter_AB_code = self.inter_AB (gpu_dst_code) gpu_dst_inter_AB_code = self.inter_AB (gpu_dst_code)
gpu_dst_code = tf.concat([gpu_dst_inter_B_code,gpu_dst_inter_AB_code],-1) gpu_dst_code = tf.concat([gpu_dst_inter_B_code,gpu_dst_inter_AB_code], nn.conv2d_ch_axis )
gpu_src_dst_code = tf.concat([gpu_dst_inter_AB_code,gpu_dst_inter_AB_code],-1) gpu_src_dst_code = tf.concat([gpu_dst_inter_AB_code,gpu_dst_inter_AB_code], nn.conv2d_ch_axis )
gpu_pred_src_src, gpu_pred_src_srcm = self.decoder(gpu_src_code) gpu_pred_src_src, gpu_pred_src_srcm = self.decoder(gpu_src_code)
gpu_pred_dst_dst, gpu_pred_dst_dstm = self.decoder(gpu_dst_code) gpu_pred_dst_dst, gpu_pred_dst_dstm = self.decoder(gpu_dst_code)
@ -503,7 +508,7 @@ class SAEHDModel(ModelBase):
gpu_src_loss = tf.reduce_mean ( 10*nn.tf_dssim(gpu_target_srcmasked_opt, gpu_pred_src_src_masked_opt, max_val=1.0, filter_size=int(resolution/11.6)), axis=[1]) gpu_src_loss = tf.reduce_mean ( 10*nn.tf_dssim(gpu_target_srcmasked_opt, gpu_pred_src_src_masked_opt, max_val=1.0, filter_size=int(resolution/11.6)), axis=[1])
gpu_src_loss += tf.reduce_mean ( 10*tf.square ( gpu_target_srcmasked_opt - gpu_pred_src_src_masked_opt ), axis=[1,2,3]) gpu_src_loss += tf.reduce_mean ( 10*tf.square ( gpu_target_srcmasked_opt - gpu_pred_src_src_masked_opt ), axis=[1,2,3])
if learn_mask: if learn_mask:
gpu_src_loss += tf.reduce_mean ( tf.square( gpu_target_srcm - gpu_pred_src_srcm ),axis=[1,2,3] ) gpu_src_loss += tf.reduce_mean ( 10*tf.square( gpu_target_srcm - gpu_pred_src_srcm ),axis=[1,2,3] )
face_style_power = self.options['face_style_power'] / 100.0 face_style_power = self.options['face_style_power'] / 100.0
if face_style_power != 0 and not self.pretrain: if face_style_power != 0 and not self.pretrain:
@ -517,14 +522,14 @@ class SAEHDModel(ModelBase):
gpu_dst_loss = tf.reduce_mean ( 10*nn.tf_dssim(gpu_target_dst_masked_opt, gpu_pred_dst_dst_masked_opt, max_val=1.0, filter_size=int(resolution/11.6) ), axis=[1]) gpu_dst_loss = tf.reduce_mean ( 10*nn.tf_dssim(gpu_target_dst_masked_opt, gpu_pred_dst_dst_masked_opt, max_val=1.0, filter_size=int(resolution/11.6) ), axis=[1])
gpu_dst_loss += tf.reduce_mean ( 10*tf.square( gpu_target_dst_masked_opt- gpu_pred_dst_dst_masked_opt ), axis=[1,2,3]) gpu_dst_loss += tf.reduce_mean ( 10*tf.square( gpu_target_dst_masked_opt- gpu_pred_dst_dst_masked_opt ), axis=[1,2,3])
if learn_mask: if learn_mask:
gpu_dst_loss += tf.reduce_mean ( tf.square( gpu_target_dstm - gpu_pred_dst_dstm ),axis=[1,2,3] ) gpu_dst_loss += tf.reduce_mean ( 10*tf.square( gpu_target_dstm - gpu_pred_dst_dstm ),axis=[1,2,3] )
gpu_src_losses += [gpu_src_loss] gpu_src_losses += [gpu_src_loss]
gpu_dst_losses += [gpu_dst_loss] gpu_dst_losses += [gpu_dst_loss]
gpu_src_dst_loss = gpu_src_loss + gpu_dst_loss gpu_src_dst_loss = gpu_src_loss + gpu_dst_loss
if self.options['true_face_training']: if self.options['true_face_power'] != 0:
def DLoss(labels,logits): def DLoss(labels,logits):
return tf.reduce_mean( tf.nn.sigmoid_cross_entropy_with_logits(labels=labels, logits=logits), axis=[1,2,3]) return tf.reduce_mean( tf.nn.sigmoid_cross_entropy_with_logits(labels=labels, logits=logits), axis=[1,2,3])
@ -534,7 +539,7 @@ class SAEHDModel(ModelBase):
gpu_dst_code_d = self.dis( gpu_dst_code ) gpu_dst_code_d = self.dis( gpu_dst_code )
gpu_dst_code_d_ones = tf.ones_like(gpu_dst_code_d) gpu_dst_code_d_ones = tf.ones_like(gpu_dst_code_d)
gpu_src_dst_loss += 0.01*DLoss(gpu_src_code_d_ones, gpu_src_code_d) gpu_src_dst_loss += self.options['true_face_power']*DLoss(gpu_src_code_d_ones, gpu_src_code_d)
gpu_D_loss = (DLoss(gpu_src_code_d_ones , gpu_dst_code_d) + \ gpu_D_loss = (DLoss(gpu_src_code_d_ones , gpu_dst_code_d) + \
DLoss(gpu_src_code_d_zeros, gpu_src_code_d) ) * 0.5 DLoss(gpu_src_code_d_zeros, gpu_src_code_d) ) * 0.5
@ -546,35 +551,20 @@ class SAEHDModel(ModelBase):
# Average losses and gradients, and create optimizer update ops # Average losses and gradients, and create optimizer update ops
with tf.device (models_opt_device): with tf.device (models_opt_device):
if gpu_count == 1: pred_src_src = nn.tf_concat(gpu_pred_src_src_list, 0)
pred_src_src = gpu_pred_src_src_list[0] pred_dst_dst = nn.tf_concat(gpu_pred_dst_dst_list, 0)
pred_dst_dst = gpu_pred_dst_dst_list[0] pred_src_dst = nn.tf_concat(gpu_pred_src_dst_list, 0)
pred_src_dst = gpu_pred_src_dst_list[0] pred_src_srcm = nn.tf_concat(gpu_pred_src_srcm_list, 0)
pred_src_srcm = gpu_pred_src_srcm_list[0] pred_dst_dstm = nn.tf_concat(gpu_pred_dst_dstm_list, 0)
pred_dst_dstm = gpu_pred_dst_dstm_list[0] pred_src_dstm = nn.tf_concat(gpu_pred_src_dstm_list, 0)
pred_src_dstm = gpu_pred_src_dstm_list[0]
src_loss = gpu_src_losses[0]
dst_loss = gpu_dst_losses[0]
src_dst_loss_gv = gpu_src_dst_loss_gvs[0]
else:
pred_src_src = tf.concat(gpu_pred_src_src_list, 0)
pred_dst_dst = tf.concat(gpu_pred_dst_dst_list, 0)
pred_src_dst = tf.concat(gpu_pred_src_dst_list, 0)
pred_src_srcm = tf.concat(gpu_pred_src_srcm_list, 0)
pred_dst_dstm = tf.concat(gpu_pred_dst_dstm_list, 0)
pred_src_dstm = tf.concat(gpu_pred_src_dstm_list, 0)
src_loss = nn.tf_average_tensor_list(gpu_src_losses) src_loss = nn.tf_average_tensor_list(gpu_src_losses)
dst_loss = nn.tf_average_tensor_list(gpu_dst_losses) dst_loss = nn.tf_average_tensor_list(gpu_dst_losses)
src_dst_loss_gv = nn.tf_average_gv_list (gpu_src_dst_loss_gvs) src_dst_loss_gv = nn.tf_average_gv_list (gpu_src_dst_loss_gvs)
if self.options['true_face_training']:
D_loss_gv = nn.tf_average_gv_list(gpu_D_loss_gvs)
src_dst_loss_gv_op = self.src_dst_opt.get_update_op (src_dst_loss_gv ) src_dst_loss_gv_op = self.src_dst_opt.get_update_op (src_dst_loss_gv )
if self.options['true_face_training']: if self.options['true_face_power'] != 0:
D_loss_gv = nn.tf_average_gv_list(gpu_D_loss_gvs)
D_loss_gv_op = self.D_opt.get_update_op (D_loss_gv ) D_loss_gv_op = self.D_opt.get_update_op (D_loss_gv )
@ -594,7 +584,7 @@ class SAEHDModel(ModelBase):
return s, d return s, d
self.src_dst_train = src_dst_train self.src_dst_train = src_dst_train
if self.options['true_face_training']: if self.options['true_face_power'] != 0:
def D_train(warped_src, warped_dst): def D_train(warped_src, warped_dst):
nn.tf_sess.run ([D_loss_gv_op], feed_dict={self.warped_src: warped_src, self.warped_dst: warped_dst}) nn.tf_sess.run ([D_loss_gv_op], feed_dict={self.warped_src: warped_src, self.warped_dst: warped_dst})
self.D_train = D_train self.D_train = D_train
@ -622,8 +612,8 @@ class SAEHDModel(ModelBase):
gpu_dst_code = self.encoder (self.warped_dst) gpu_dst_code = self.encoder (self.warped_dst)
gpu_dst_inter_B_code = self.inter_B (gpu_dst_code) gpu_dst_inter_B_code = self.inter_B (gpu_dst_code)
gpu_dst_inter_AB_code = self.inter_AB (gpu_dst_code) gpu_dst_inter_AB_code = self.inter_AB (gpu_dst_code)
gpu_dst_code = tf.concat([gpu_dst_inter_B_code,gpu_dst_inter_AB_code],-1) gpu_dst_code = tf.concat([gpu_dst_inter_B_code,gpu_dst_inter_AB_code], nn.conv2d_ch_axis)
gpu_src_dst_code = tf.concat([gpu_dst_inter_AB_code,gpu_dst_inter_AB_code],-1) gpu_src_dst_code = tf.concat([gpu_dst_inter_AB_code,gpu_dst_inter_AB_code], nn.conv2d_ch_axis)
gpu_pred_src_dst, gpu_pred_src_dstm = self.decoder(gpu_src_dst_code) gpu_pred_src_dst, gpu_pred_src_dstm = self.decoder(gpu_src_dst_code)
_, gpu_pred_dst_dstm = self.decoder(gpu_dst_code) _, gpu_pred_dst_dstm = self.decoder(gpu_dst_code)
@ -683,16 +673,16 @@ class SAEHDModel(ModelBase):
self.set_training_data_generators ([ self.set_training_data_generators ([
SampleGeneratorFace(training_data_src_path, random_ct_samples_path=random_ct_samples_path, debug=self.is_debug(), batch_size=self.get_batch_size(), SampleGeneratorFace(training_data_src_path, random_ct_samples_path=random_ct_samples_path, debug=self.is_debug(), batch_size=self.get_batch_size(),
sample_process_options=SampleProcessor.Options(random_flip=self.random_flip), sample_process_options=SampleProcessor.Options(random_flip=self.random_flip),
output_sample_types = [ {'types' : (t_img_warped, face_type, t.MODE_BGR), 'resolution':resolution, 'ct_mode': self.options['ct_mode'] }, output_sample_types = [ {'types' : (t_img_warped, face_type, t.MODE_BGR), 'data_format':nn.data_format, 'resolution': resolution, 'ct_mode': self.options['ct_mode'] },
{'types' : (t.IMG_TRANSFORMED, face_type, t.MODE_BGR), 'resolution': resolution, 'ct_mode': self.options['ct_mode'] }, {'types' : (t.IMG_TRANSFORMED, face_type, t.MODE_BGR), 'data_format':nn.data_format, 'resolution': resolution, 'ct_mode': self.options['ct_mode'] },
{'types' : (t.IMG_TRANSFORMED, face_type, t.MODE_M), 'resolution': resolution } ], {'types' : (t.IMG_TRANSFORMED, face_type, t.MODE_M), 'data_format':nn.data_format, 'resolution': resolution } ],
generators_count=src_generators_count ), generators_count=src_generators_count ),
SampleGeneratorFace(training_data_dst_path, debug=self.is_debug(), batch_size=self.get_batch_size(), SampleGeneratorFace(training_data_dst_path, debug=self.is_debug(), batch_size=self.get_batch_size(),
sample_process_options=SampleProcessor.Options(random_flip=self.random_flip), sample_process_options=SampleProcessor.Options(random_flip=self.random_flip),
output_sample_types = [ {'types' : (t_img_warped, face_type, t.MODE_BGR), 'resolution':resolution}, output_sample_types = [ {'types' : (t_img_warped, face_type, t.MODE_BGR), 'data_format':nn.data_format, 'resolution': resolution},
{'types' : (t.IMG_TRANSFORMED, face_type, t.MODE_BGR), 'resolution': resolution}, {'types' : (t.IMG_TRANSFORMED, face_type, t.MODE_BGR), 'data_format':nn.data_format, 'resolution': resolution},
{'types' : (t.IMG_TRANSFORMED, face_type, t.MODE_M), 'resolution': resolution} ], {'types' : (t.IMG_TRANSFORMED, face_type, t.MODE_M), 'data_format':nn.data_format, 'resolution': resolution} ],
generators_count=dst_generators_count ) generators_count=dst_generators_count )
]) ])
@ -713,7 +703,7 @@ class SAEHDModel(ModelBase):
src_loss, dst_loss = self.src_dst_train (warped_src, target_src, target_srcm, warped_dst, target_dst, target_dstm) src_loss, dst_loss = self.src_dst_train (warped_src, target_src, target_srcm, warped_dst, target_dst, target_dstm)
if self.options['true_face_training'] and not self.pretrain: if self.options['true_face_power'] != 0 and not self.pretrain:
self.D_train (warped_src, warped_dst) self.D_train (warped_src, warped_dst)
return ( ('src_loss', src_loss), ('dst_loss', dst_loss), ) return ( ('src_loss', src_loss), ('dst_loss', dst_loss), )
@ -728,10 +718,12 @@ class SAEHDModel(ModelBase):
for sample_list in samples ] for sample_list in samples ]
if self.options['learn_mask']: if self.options['learn_mask']:
S, D, SS, DD, DDM, SD, SDM = [ np.clip(x, 0.0, 1.0) for x in ([target_src,target_dst] + self.AE_view (target_src, target_dst) ) ] S, D, SS, DD, DDM, SD, SDM = [ np.clip( nn.to_data_format(x,"NHWC", self.model_data_format), 0.0, 1.0) for x in ([target_src,target_dst] + self.AE_view (target_src, target_dst) ) ]
DDM, SDM, = [ np.repeat (x, (3,), -1) for x in [DDM, SDM] ] DDM, SDM, = [ np.repeat (x, (3,), -1) for x in [DDM, SDM] ]
else: else:
S, D, SS, DD, SD, = [ np.clip(x, 0.0, 1.0) for x in ([target_src,target_dst] + self.AE_view (target_src, target_dst) ) ] S, D, SS, DD, SD, = [ np.clip( nn.to_data_format(x,"NHWC", self.model_data_format) , 0.0, 1.0) for x in ([target_src,target_dst] + self.AE_view (target_src, target_dst) ) ]
target_srcm, target_dstm = [ nn.to_data_format(x,"NHWC", self.model_data_format) for x in ([target_srcm, target_dstm] )]
result = [] result = []
st = [] st = []
@ -753,12 +745,16 @@ class SAEHDModel(ModelBase):
return result return result
def predictor_func (self, face=None): def predictor_func (self, face=None):
face = face[None,...]
face = nn.to_data_format(face, self.model_data_format, "NHWC")
if self.options['learn_mask']: if self.options['learn_mask']:
bgr, mask_dst_dstm, mask_src_dstm = self.AE_merge (face[np.newaxis,...]) bgr, mask_dst_dstm, mask_src_dstm = [ nn.to_data_format(x,"NHWC", self.model_data_format).astype(np.float32) for x in self.AE_merge (face) ]
mask = mask_dst_dstm[0] * mask_src_dstm[0] mask = mask_dst_dstm[0] * mask_src_dstm[0]
return bgr[0], mask[...,0] return bgr[0], mask[...,0]
else: else:
bgr, = self.AE_merge (face[np.newaxis,...]) bgr, = [ nn.to_data_format(x,"NHWC", self.model_data_format).astype(np.float32) for x in self.AE_merge (face) ]
return bgr[0] return bgr[0]
#override #override

2
requirements-colab.txt
View file

@ -6,4 +6,4 @@ ffmpeg-python==0.1.17
scikit-image==0.14.2 scikit-image==0.14.2
scipy==1.4.1 scipy==1.4.1
colorama colorama
tensorflow-gpu==1.13.1 tensorflow-gpu==1.13.2

2
requirements-cuda.txt
View file

@ -6,4 +6,4 @@ ffmpeg-python==0.1.17
scikit-image==0.14.2 scikit-image==0.14.2
scipy==1.4.1 scipy==1.4.1
colorama colorama
tensorflow-gpu==1.12.0 tensorflow-gpu==1.13.2

4
samplelib/SampleHost.py
View file

@ -15,10 +15,6 @@ from .Sample import Sample, SampleType
class SampleHost: class SampleHost:
samples_cache = dict() samples_cache = dict()
@staticmethod @staticmethod
def get_person_id_max_count(samples_path): def get_person_id_max_count(samples_path):

6
samplelib/SampleProcessor.py
View file

@ -126,6 +126,8 @@ class SampleProcessor(object):
ct_mode = opts.get('ct_mode', 'None') ct_mode = opts.get('ct_mode', 'None')
normalize_tanh = opts.get('normalize_tanh', False) normalize_tanh = opts.get('normalize_tanh', False)
data_format = opts.get('data_format', 'NHWC')
img_type = SPTF.NONE img_type = SPTF.NONE
target_face_type = SPTF.NONE target_face_type = SPTF.NONE
@ -300,6 +302,10 @@ class SampleProcessor(object):
else: else:
img = np.clip (img, 0.0, 1.0) img = np.clip (img, 0.0, 1.0)
if data_format == "NCHW":
img = np.transpose(img, (2,0,1) )
outputs_sample.append ( img ) outputs_sample.append ( img )
outputs += [outputs_sample] outputs += [outputs_sample]