github/DeepFaceLab
1
0
Fork 0
mirror of https://github.com/iperov/DeepFaceLab.git synced 2025-07-06 04:52:13 -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 ca (init_ops.Initializer):
def __init__(self, dtype=None):
pass
def __call__(self, shape, dtype=None, partition_info=None):
return tf.zeros( shape, name="_cai_")
return tf.zeros( shape, dtype=dtype, name="_cai_")
@staticmethod
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:])
w_val = d.get(sub_w_name, None)
w_val = np.reshape( w_val, w.shape.as_list() )
if w_val is None:
io.log_err(f"Weight {w.name} was not loaded from file {filename}")
tuples.append ( (w, w.initializer) )
@ -255,13 +257,22 @@ def initialize_layers(nn):
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):
if not isinstance(shapes_list, list):
raise ValueError("shapes_list must be a list.")
self.run_placeholders = []
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)
@ -292,6 +303,9 @@ def initialize_layers(nn):
if not isinstance(dilations, int):
raise ValueError ("dilations must be an int type")
if dtype is None:
dtype = nn.tf_floatx
if isinstance(padding, str):
if padding == "SAME":
padding = ( (kernel_size - 1) * dilations + 1 ) // 2
@ -302,36 +316,47 @@ def initialize_layers(nn):
if isinstance(padding, int):
if padding != 0:
if nn.data_format == "NHWC":
padding = [ [0,0], [padding,padding], [padding,padding], [0,0] ]
else:
padding = [ [0,0], [0,0], [padding,padding], [padding,padding] ]
else:
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.out_ch = out_ch
self.kernel_size = kernel_size
self.strides = [1,strides,strides,1]
self.strides = strides
self.padding = padding
self.dilations = [1,dilations,dilations,1]
self.dilations = dilations
self.use_bias = use_bias
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.trainable = trainable
if dtype is None:
dtype = nn.tf_floatx
self.dtype = dtype
super().__init__(**kwargs)
def build_weights(self):
kernel_initializer = self.kernel_initializer
if kernel_initializer is None:
if self.use_wscale:
gain = 1.0 if self.kernel_size == 1 else np.sqrt(2)
fan_in = self.kernel_size*self.kernel_size*self.in_ch
he_std = gain / np.sqrt(fan_in) # He init
self.wscale = tf.constant(he_std, 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)
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:
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):
weights = [self.weight]
@ -357,9 +382,13 @@ def initialize_layers(nn):
if self.padding is not None:
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:
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
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 ):
if not isinstance(strides, int):
raise ValueError ("strides must be an int type")
if dtype is None:
dtype = nn.tf_floatx
self.in_ch = in_ch
self.out_ch = out_ch
self.kernel_size = kernel_size
@ -383,33 +416,30 @@ def initialize_layers(nn):
self.padding = padding
self.use_bias = use_bias
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.trainable = trainable
if dtype is None:
dtype = nn.tf_floatx
self.dtype = dtype
super().__init__(**kwargs)
def build_weights(self):
kernel_initializer = self.kernel_initializer
if kernel_initializer is None:
if self.use_wscale:
gain = 1.0 if self.kernel_size == 1 else np.sqrt(2)
fan_in = self.kernel_size*self.kernel_size*self.in_ch
he_std = gain / np.sqrt(fan_in) # He init
self.wscale = tf.constant(he_std, 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)
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:
bias_initializer = self.bias_initializer
if bias_initializer is None:
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):
weights = [self.weight]
@ -420,21 +450,34 @@ def initialize_layers(nn):
def __call__(self, x):
shape = x.shape
if nn.data_format == "NHWC":
h,w,c = shape[1], shape[2], shape[3]
output_shape = tf.stack ( (tf.shape(x)[0],
self.deconv_length(w, self.strides, self.kernel_size, self.padding),
self.deconv_length(h, self.strides, self.kernel_size, self.padding),
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
if self.use_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:
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
def __str__(self):
@ -459,10 +502,13 @@ def initialize_layers(nn):
def __init__(self, filt_size=3, stride=2, **kwargs ):
self.strides = [1,stride,stride,1]
self.filt_size = filt_size
self.padding = [ [0,0],
[ 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)) ],
[0,0] ]
pad = [ int(1.*(filt_size-1)/2), int(np.ceil(1.*(filt_size-1)/2)) ]
if nn.data_format == "NHWC":
self.padding = [ [0,0], pad, pad, [0,0] ]
else:
self.padding = [ [0,0], [0,0], pad, pad ]
if(self.filt_size==1):
a = np.array([1.,])
elif(self.filt_size==2):
@ -512,7 +558,8 @@ def initialize_layers(nn):
self.bias_initializer = bias_initializer
self.trainable = trainable
if dtype is None:
dtype = tf.float32
dtype = nn.tf_floatx
self.dtype = dtype
super().__init__(**kwargs)
@ -523,14 +570,15 @@ def initialize_layers(nn):
weight_shape = (self.in_ch,self.out_ch)
kernel_initializer = self.kernel_initializer
if kernel_initializer is None:
if self.use_wscale:
gain = 1.0
fan_in = np.prod( weight_shape[:-1] )
he_std = gain / np.sqrt(fan_in) # He init
self.wscale = tf.constant(he_std, 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)
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
if bias_initializer is None:
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):
weights = [self.weight]
@ -559,7 +607,7 @@ def initialize_layers(nn):
x = tf.reduce_max(x, axis=-1)
if self.use_bias:
x = x + self.bias
x = tf.add(x, tf.reshape(self.bias, (1,self.out_ch) ) )
return x
nn.Dense = Dense
@ -568,31 +616,38 @@ def initialize_layers(nn):
"""
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.eps = eps
self.momentum = momentum
if dtype is None:
dtype = nn.tf_floatx
self.dtype = dtype
self.shape = (1,1,1,dim)
super().__init__(**kwargs)
def build_weights(self):
self.weight = tf.get_variable("weight", self.shape, dtype=self.dtype, initializer=tf.initializers.ones() )
self.bias = tf.get_variable("bias", self.shape, 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_var = tf.get_variable("running_var", self.shape, dtype=self.dtype, initializer=tf.initializers.zeros(), trainable=False )
self.weight = tf.get_variable("weight", (self.dim,), dtype=self.dtype, initializer=tf.initializers.ones() )
self.bias = tf.get_variable("bias", (self.dim,), dtype=self.dtype, initializer=tf.initializers.zeros() )
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.dim,), dtype=self.dtype, initializer=tf.initializers.zeros(), trainable=False )
def get_weights(self):
return [self.weight, self.bias, self.running_mean, self.running_var]
def __call__(self, x):
x = (x - self.running_mean) / tf.sqrt( self.running_var + self.eps )
x *= self.weight
x += self.bias
if nn.data_format == "NHWC":
shape = (1,1,1,self.dim)
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
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:
1) change env variables based on DeviceConfig before import tensorflow
2) multiprocesses will import tensorflow every spawn
NCHW speed up training for 10-20%.
"""
import os
import sys
from pathlib import Path
import numpy as np
from core.interact import interact as io
from .device import Devices
class nn():
current_DeviceConfig = None
tf = None
tf_sess = 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
tf_get_value = None
@ -34,17 +49,18 @@ class nn():
tf_gradients = None
tf_average_gv_list = None
tf_average_tensor_list = None
tf_dot = None
tf_concat = None
tf_gelu = None
tf_upsample2d = None
tf_upsample2d_bilinear = None
tf_flatten = None
tf_reshape_4D = None
tf_random_binomial = None
tf_gaussian_blur = None
tf_style_loss = None
tf_channel_histogram = None
tf_histogram = None
tf_dssim = None
tf_space_to_depth = None
tf_depth_to_space = None
# Layers
Saveable = None
@ -64,7 +80,8 @@ class nn():
TFRMSpropOptimizer = None
@staticmethod
def initialize(device_config=None):
def initialize(device_config=None, floatx="float32", data_format="NHWC"):
if nn.tf is None:
if device_config is None:
device_config = nn.getCurrentDeviceConfig()
@ -74,11 +91,8 @@ class nn():
if 'CUDA_VISIBLE_DEVICES' in os.environ.keys():
os.environ.pop('CUDA_VISIBLE_DEVICES')
os.environ['CUDA_​CACHE_​MAXSIZE'] = '536870912' #512Mb (32mb default)
first_run = False
if not device_config.cpu_only:
if len(device_config.devices) != 0:
if sys.platform[0:3] == 'win':
if all( [ x.name == device_config.devices[0].name for x in device_config.devices ] ):
devices_str = "_" + device_config.devices[0].name.replace(' ','_')
@ -92,6 +106,7 @@ class nn():
first_run = True
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_CPP_MIN_LOG_LEVEL'] = '2' # tf log errors only
@ -102,11 +117,16 @@ class nn():
io.log_info("Caching GPU kernels...")
import tensorflow as tf
import logging
logging.getLogger('tensorflow').setLevel(logging.ERROR)
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})
else:
nn.tf_default_device = "/GPU:0"
config = tf.ConfigProto()
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
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 .layers import initialize_layers
from .initializers import initialize_initializers
@ -130,10 +147,68 @@ class nn():
if nn.tf_sess is None:
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
def 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
def getCurrentDeviceConfig():
if nn.current_DeviceConfig is None:
@ -159,6 +234,13 @@ class nn():
nn.tf_sess.close()
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
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
if e: e.__enter__()
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) ]
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)
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) } )
if e: e.__exit__(None, None, None)
def get_update_op(self, grads_vars):
updates = []
lr = self.lr
if self.clipnorm > 0.0:
norm = tf.sqrt( sum([tf.reduce_sum(tf.square(g)) for g,v in grads_vars]))
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)
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:
lr_rnd = self.lr_rnds_dict[v.name]
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
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
if e is not None: e.__enter__()
result = []
@ -58,6 +61,9 @@ def initialize_tensor_ops(nn):
nn.tf_average_gv_list = tf_average_gv_list
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
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)
@ -65,36 +71,14 @@ def initialize_tensor_ops(nn):
return result
nn.tf_average_tensor_list = tf_average_tensor_list
def tf_dot(x, y):
if x.shape.ndims > 2 or y.shape.ndims > 2:
x_shape = []
for i, s in zip( x.shape.as_list(), array_ops.unstack(array_ops.shape(x))):
if i is not None:
x_shape.append(i)
else:
x_shape.append(s)
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_concat (tensors_list, axis):
"""
Better version.
"""
if len(tensors_list) == 1:
return tensors_list[0]
return tf.concat(tensors_list, axis)
nn.tf_concat = tf_concat
def tf_gelu(x):
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
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) )
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) )
nn.tf_upsample2d_bilinear = tf_upsample2d_bilinear
def tf_flatten(x, dynamic_dims=False):
"""
dynamic_dims allows to flatten without knowing size on input dims
"""
if dynamic_dims:
sh = tf.shape(x)
return tf.reshape (x, (sh[0], tf.reduce_prod(sh[1:]) ) )
else:
def tf_flatten(x):
if nn.data_format == "NHWC":
# match NCHW version in order to switch data_format without problems
x = tf.transpose(x, (0,3,1,2) )
return tf.reshape (x, (-1, np.prod(x.shape[1:])) )
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):
if dtype is None:
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)])
np_kernel = np.outer(kernel_1d, kernel_1d).astype(np.float32)
kernel = np_kernel / np.sum(np_kernel)
return kernel
return kernel, kernel_size
gauss_kernel = make_kernel(radius)
gauss_kernel = gauss_kernel[:, :,np.newaxis, np.newaxis]
kernel_size = gauss_kernel.shape[0]
inputs = [ input[:,:,:,i:i+1] for i in range( input.shape[-1] ) ]
gauss_kernel, kernel_size = make_kernel(radius)
padding = kernel_size//2
if padding != 0:
if nn.data_format == "NHWC":
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 = []
for i in range(len(inputs)):
x = inputs[i]
if kernel_size != 0:
padding = kernel_size//2
x = tf.pad (x, [ [0,0], [padding,padding], [padding,padding], [0,0] ] )
for i in range(input.shape[nn.conv2d_ch_axis]):
x = input[:,:,:,i:i+1] if nn.data_format == "NHWC" \
else input[:,i:i+1,:,:]
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
def tf_style_loss(target, style, gaussian_blur_radius=0.0, loss_weight=1.0, step_size=1):
def sd(content, style, loss_weight):
content_nc = content.shape[-1]
style_nc = style.shape[-1]
content_nc = content.shape[ nn.conv2d_ch_axis ]
style_nc = style.shape[nn.conv2d_ch_axis]
if content_nc != style_nc:
raise Exception("style_loss() content_nc != style_nc")
axes = [1,2]
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_mean, c_var = tf.nn.moments(content, axes=nn.conv2d_spatial_axes, keep_dims=True)
s_mean, s_var = tf.nn.moments(style, axes=nn.conv2d_spatial_axes, keep_dims=True)
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])
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 )
if gaussian_blur_radius > 0.0:
@ -187,46 +185,29 @@ def initialize_tensor_ops(nn):
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):
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):
#Function to mimic the 'fspecial' gaussian MATLAB function.
coords = np.arange(0, size, dtype=nn.np_floatx)
coords -= (size - 1 ) / 2.0
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
if not_float32:
img_dtype = img1.dtype
img1 = tf.cast(img1, tf.float32)
img2 = tf.cast(img2, tf.float32)
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):
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
c2 = (k2 * max_val) ** 2
@ -242,10 +223,44 @@ def initialize_tensor_ops(nn):
c2 *= 1.0 #compensation factor
cs = (num1 - num0 + c2) / (den1 - den0 + c2)
ssim_val = tf.reduce_mean(luminance * cs, axis=(-3, -2) )
return(1.0 - ssim_val ) / 2.0
ssim_val = tf.reduce_mean(luminance * cs, axis=nn.conv2d_spatial_axes )
dssim = (1.0 - ssim_val ) / 2.0
if not_float32:
dssim = tf.cast(dssim, img_dtype)
return 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):
srgb_pixels = tf.reshape(srgb, [-1, 3])
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():
raise Exception("Unable to load FANExtractor model")
nn.initialize()
nn.initialize(data_format="NHWC")
tf = nn.tf
class ConvBlock(nn.ModelBase):
@ -29,10 +29,10 @@ class FANExtractor(object):
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.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.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 )
if self.in_planes != self.out_planes:
@ -55,6 +55,7 @@ class FANExtractor(object):
x = self.bn3(x)
x = tf.nn.relu(x)
x = out3 = self.conv3(x)
x = tf.concat ([out1, out2, out3], axis=-1)
if self.in_planes != self.out_planes:
@ -148,7 +149,9 @@ class FANExtractor(object):
if i < 4 - 1:
ll = self.bl[i](ll)
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
if place_model_on_cpu:
@ -159,7 +162,7 @@ class FANExtractor(object):
self.model.load_weights(str(model_path))
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):
if len(rects) == 0:
@ -197,7 +200,7 @@ class FANExtractor(object):
predicted = []
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)

8
facelib/FaceEnhancer.py
View file

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

4
facelib/S3FDExtractor.py
View file

@ -8,7 +8,7 @@ from core.leras import nn
class S3FDExtractor(object):
def __init__(self, place_model_on_cpu=False):
nn.initialize()
nn.initialize(data_format="NHWC")
tf = nn.tf
model_path = Path(__file__).parent / "S3FD.npy"
@ -167,7 +167,7 @@ class S3FDExtractor(object):
self.model.load_weights (model_path)
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):
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):
VERSION = 1
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
class Ternaus(nn.ModelBase):
@ -125,7 +125,7 @@ class TernausNet(object):
self.net.init_weights()
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:
raise Exception("training not supported yet")
@ -195,124 +195,3 @@ class TernausNet(object):
result = result[0]
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
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
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
#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.
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
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_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'=={" "*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
else:
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):
#override
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
conv_kernel_initializer = nn.initializers.ca
conv_kernel_initializer = nn.initializers.ca()
class Downscale(nn.ModelBase):
def __init__(self, in_ch, out_ch, kernel_size=5, dilations=1, subpixel=True, use_activator=True, *kwargs ):
@ -39,7 +41,7 @@ class QModel(ModelBase):
x = self.conv1(x)
if self.subpixel:
x = tf.nn.space_to_depth(x, 2)
x = nn.tf_space_to_depth(x, 2)
if self.use_activator:
x = nn.tf_gelu(x)
@ -71,7 +73,7 @@ class QModel(ModelBase):
def forward(self, x):
x = self.conv1(x)
x = nn.tf_gelu(x)
x = tf.nn.depth_to_space(x, 2)
x = nn.tf_depth_to_space(x, 2)
return x
class ResidualBlock(nn.ModelBase):
@ -109,7 +111,7 @@ class QModel(ModelBase):
def forward(self, inp):
x = self.dense1(inp)
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.res1(x)
return x
@ -167,8 +169,8 @@ class QModel(ModelBase):
input_nc = 3
output_nc = 3
bgr_shape = (resolution, resolution, output_nc)
mask_shape = (resolution, resolution, 1)
bgr_shape = nn.get4Dshape(resolution,resolution,input_nc)
mask_shape = nn.get4Dshape(resolution,resolution,1)
lowest_dense_res = resolution // 16
self.model_filename_list = []
@ -176,22 +178,22 @@ class QModel(ModelBase):
with tf.device ('/CPU:0'):
#Place holders on CPU
self.warped_src = tf.placeholder (tf.float32, (None,)+bgr_shape)
self.warped_dst = tf.placeholder (tf.float32, (None,)+bgr_shape)
self.warped_src = tf.placeholder (nn.tf_floatx, bgr_shape)
self.warped_dst = tf.placeholder (nn.tf_floatx, bgr_shape)
self.target_src = tf.placeholder (tf.float32, (None,)+bgr_shape)
self.target_dst = tf.placeholder (tf.float32, (None,)+bgr_shape)
self.target_src = tf.placeholder (nn.tf_floatx, bgr_shape)
self.target_dst = tf.placeholder (nn.tf_floatx, bgr_shape)
self.target_srcm = tf.placeholder (tf.float32, (None,)+mask_shape)
self.target_dstm = tf.placeholder (tf.float32, (None,)+mask_shape)
self.target_srcm = tf.placeholder (nn.tf_floatx, mask_shape)
self.target_dstm = tf.placeholder (nn.tf_floatx, mask_shape)
# Initializing model classes
with tf.device (models_opt_device):
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')
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_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*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*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_dst_losses += [gpu_dst_loss]
@ -286,29 +288,16 @@ class QModel(ModelBase):
# Average losses and gradients, and create optimizer update ops
with tf.device (models_opt_device):
if gpu_count == 1:
pred_src_src = gpu_pred_src_src_list[0]
pred_dst_dst = gpu_pred_dst_dst_list[0]
pred_src_dst = gpu_pred_src_dst_list[0]
pred_src_srcm = gpu_pred_src_srcm_list[0]
pred_dst_dstm = gpu_pred_dst_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)
pred_src_src = nn.tf_concat(gpu_pred_src_src_list, 0)
pred_dst_dst = nn.tf_concat(gpu_pred_dst_dst_list, 0)
pred_src_dst = nn.tf_concat(gpu_pred_src_dst_list, 0)
pred_src_srcm = nn.tf_concat(gpu_pred_src_srcm_list, 0)
pred_dst_dstm = nn.tf_concat(gpu_pred_dst_dstm_list, 0)
pred_src_dstm = nn.tf_concat(gpu_pred_src_dstm_list, 0)
src_loss = nn.tf_average_tensor_list(gpu_src_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_op = self.src_dst_opt.get_update_op (src_dst_loss_gv)
# Initializing training and view functions
@ -341,13 +330,11 @@ class QModel(ModelBase):
_, gpu_pred_dst_dstm = self.decoder_dst(gpu_dst_code)
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})
self.AE_merge = AE_merge
# Loading/initializing all models/optimizers weights
for model, filename in io.progress_bar_generator(self.model_filename_list, "Initializing models"):
do_init = self.is_first_run()
@ -368,7 +355,6 @@ class QModel(ModelBase):
model.init_weights()
# initializing sample generators
if self.is_training:
t = SampleProcessor.Types
face_type = t.FACE_TYPE_FULL
@ -384,16 +370,16 @@ class QModel(ModelBase):
self.set_training_data_generators ([
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),
output_sample_types = [ {'types' : (t.IMG_WARPED_TRANSFORMED, face_type, t.MODE_BGR), 'resolution':resolution, },
{'types' : (t.IMG_TRANSFORMED, face_type, t.MODE_BGR), 'resolution': resolution, },
{'types' : (t.IMG_TRANSFORMED, face_type, t.MODE_M), '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), 'data_format':nn.data_format, 'resolution': resolution, },
{'types' : (t.IMG_TRANSFORMED, face_type, t.MODE_M), 'data_format':nn.data_format, 'resolution': resolution } ],
generators_count=src_generators_count ),
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),
output_sample_types = [ {'types' : (t.IMG_WARPED_TRANSFORMED, face_type, t.MODE_BGR), 'resolution':resolution},
{'types' : (t.IMG_TRANSFORMED, face_type, t.MODE_BGR), 'resolution': resolution},
{'types' : (t.IMG_TRANSFORMED, face_type, t.MODE_M), '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), 'data_format':nn.data_format, 'resolution': resolution},
{'types' : (t.IMG_TRANSFORMED, face_type, t.MODE_M), 'data_format':nn.data_format, 'resolution': resolution} ],
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):
model.save_weights ( self.get_strpath_storage_for_file(filename) )
#override
def onTrainOneIter(self):
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 ]
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] ]
target_srcm, target_dstm = [ nn.to_data_format(x,"NHWC", self.model_data_format) for x in ([target_srcm, target_dstm] )]
result = []
st = []
for i in range(n_samples):
@ -456,8 +443,10 @@ class QModel(ModelBase):
return result
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]
return bgr[0], mask[...,0]

166
models/Model_SAEHD/Model.py
View file

@ -26,14 +26,6 @@ class SAEHDModel(ModelBase):
suggest_batch_size = 4
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_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 = 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_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_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_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_clipgrad = self.options['clipgrad'] = self.load_or_def_option('clipgrad', 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():
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)
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()
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['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:
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['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']:
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:
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['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:
self.set_iter(1)
#override
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
conv_kernel_initializer = nn.initializers.ca
conv_kernel_initializer = nn.initializers.ca()
class Downscale(nn.ModelBase):
def __init__(self, in_ch, out_ch, kernel_size=5, dilations=1, subpixel=True, use_activator=True, *kwargs ):
@ -125,13 +132,13 @@ class SAEHDModel(ModelBase):
self.out_ch // (4 if self.subpixel else 1),
kernel_size=self.kernel_size,
strides=1 if self.subpixel else 2,
padding='SAME', dilations=self.dilations, kernel_initializer=conv_kernel_initializer )
padding='SAME', dilations=self.dilations, kernel_initializer=conv_kernel_initializer)
def forward(self, x):
x = self.conv1(x)
if self.subpixel:
x = tf.nn.space_to_depth(x, 2)
x = nn.tf_space_to_depth(x, 2)
if self.use_activator:
x = tf.nn.leaky_relu(x, 0.1)
@ -163,7 +170,7 @@ class SAEHDModel(ModelBase):
def forward(self, x):
x = self.conv1(x)
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
class ResidualBlock(nn.ModelBase):
@ -211,7 +218,6 @@ class SAEHDModel(ModelBase):
nn.tf_flatten(self.down4(inp)) ], -1 )
else:
x = nn.tf_flatten(self.down1(inp))
return x
class Inter(nn.ModelBase):
@ -222,14 +228,14 @@ class SAEHDModel(ModelBase):
def on_build(self):
in_ch, lowest_dense_res, ae_ch, ae_out_ch = self.in_ch, self.lowest_dense_res, self.ae_ch, self.ae_out_ch
self.dense1 = nn.Dense( in_ch, ae_ch, kernel_initializer=tf.initializers.orthogonal )
self.dense2 = nn.Dense( ae_ch, lowest_dense_res * lowest_dense_res * ae_out_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 )
self.upscale1 = Upscale(ae_out_ch, ae_out_ch)
def forward(self, inp):
x = self.dense1(inp)
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)
return x
@ -307,7 +313,7 @@ class SAEHDModel(ModelBase):
class CodeDiscriminator(nn.ModelBase):
def on_build(self, in_ch, code_res, ch=256):
n_downscales = 2 + code_res // 8
n_downscales = 1 + code_res // 8
self.convs = []
prev_ch = in_ch
@ -343,8 +349,8 @@ class SAEHDModel(ModelBase):
input_nc = 3
output_nc = 3
bgr_shape = (resolution, resolution, output_nc)
mask_shape = (resolution, resolution, 1)
bgr_shape = nn.get4Dshape(resolution,resolution,input_nc)
mask_shape = nn.get4Dshape(resolution,resolution,1)
lowest_dense_res = resolution // 16
self.model_filename_list = []
@ -352,23 +358,23 @@ class SAEHDModel(ModelBase):
with tf.device ('/CPU:0'):
#Place holders on CPU
self.warped_src = tf.placeholder (tf.float32, (None,)+bgr_shape)
self.warped_dst = tf.placeholder (tf.float32, (None,)+bgr_shape)
self.warped_src = tf.placeholder (nn.tf_floatx, bgr_shape)
self.warped_dst = tf.placeholder (nn.tf_floatx, bgr_shape)
self.target_src = tf.placeholder (tf.float32, (None,)+bgr_shape)
self.target_dst = tf.placeholder (tf.float32, (None,)+bgr_shape)
self.target_src = tf.placeholder (nn.tf_floatx, bgr_shape)
self.target_dst = tf.placeholder (nn.tf_floatx, bgr_shape)
self.target_srcm = tf.placeholder (tf.float32, (None,)+mask_shape)
self.target_dstm = tf.placeholder (tf.float32, (None,)+mask_shape)
self.target_srcm = tf.placeholder (nn.tf_floatx, mask_shape)
self.target_dstm = tf.placeholder (nn.tf_floatx, mask_shape)
# Initializing model classes
with tf.device (models_opt_device):
if 'df' in archi:
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')
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_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'] ]
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.model_filename_list += [ [self.dis, 'dis.npy'] ]
elif 'liae' in archi:
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_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_B_out_ch = self.inter_B.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_channels ( (nn.tf_floatx, (None,encoder_out_ch)))
inters_out_ch = inter_AB_out_ch+inter_B_out_ch
self.decoder = Decoder(in_ch=inters_out_ch, d_ch=d_dims, d_mask_ch=d_mask_dims, is_hd='hd' in archi, name='decoder')
self.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)
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.initialize_variables ( self.dis.get_weights(), vars_on_cpu=optimizer_vars_on_cpu)
self.model_filename_list += [ (self.D_opt, 'D_opt.npy') ]
@ -466,12 +471,12 @@ class SAEHDModel(ModelBase):
elif 'liae' in archi:
gpu_src_code = self.encoder (gpu_warped_src)
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_inter_B_code = self.inter_B (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_src_dst_code = tf.concat([gpu_dst_inter_AB_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], nn.conv2d_ch_axis )
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)
@ -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*tf.square ( gpu_target_srcmasked_opt - gpu_pred_src_src_masked_opt ), axis=[1,2,3])
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
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*tf.square( gpu_target_dst_masked_opt- gpu_pred_dst_dst_masked_opt ), axis=[1,2,3])
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_dst_losses += [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):
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_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) + \
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
with tf.device (models_opt_device):
if gpu_count == 1:
pred_src_src = gpu_pred_src_src_list[0]
pred_dst_dst = gpu_pred_dst_dst_list[0]
pred_src_dst = gpu_pred_src_dst_list[0]
pred_src_srcm = gpu_pred_src_srcm_list[0]
pred_dst_dstm = gpu_pred_dst_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)
pred_src_src = nn.tf_concat(gpu_pred_src_src_list, 0)
pred_dst_dst = nn.tf_concat(gpu_pred_dst_dst_list, 0)
pred_src_dst = nn.tf_concat(gpu_pred_src_dst_list, 0)
pred_src_srcm = nn.tf_concat(gpu_pred_src_srcm_list, 0)
pred_dst_dstm = nn.tf_concat(gpu_pred_dst_dstm_list, 0)
pred_src_dstm = nn.tf_concat(gpu_pred_src_dstm_list, 0)
src_loss = nn.tf_average_tensor_list(gpu_src_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)
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 )
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 )
@ -594,7 +584,7 @@ class SAEHDModel(ModelBase):
return s, d
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):
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
@ -622,8 +612,8 @@ class SAEHDModel(ModelBase):
gpu_dst_code = self.encoder (self.warped_dst)
gpu_dst_inter_B_code = self.inter_B (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_src_dst_code = tf.concat([gpu_dst_inter_AB_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], nn.conv2d_ch_axis)
gpu_pred_src_dst, gpu_pred_src_dstm = self.decoder(gpu_src_dst_code)
_, gpu_pred_dst_dstm = self.decoder(gpu_dst_code)
@ -683,16 +673,16 @@ class SAEHDModel(ModelBase):
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(),
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'] },
{'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_M), 'resolution': resolution } ],
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), 'data_format':nn.data_format, 'resolution': resolution, 'ct_mode': self.options['ct_mode'] },
{'types' : (t.IMG_TRANSFORMED, face_type, t.MODE_M), 'data_format':nn.data_format, 'resolution': resolution } ],
generators_count=src_generators_count ),
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),
output_sample_types = [ {'types' : (t_img_warped, face_type, t.MODE_BGR), 'resolution':resolution},
{'types' : (t.IMG_TRANSFORMED, face_type, t.MODE_BGR), 'resolution': resolution},
{'types' : (t.IMG_TRANSFORMED, face_type, t.MODE_M), '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), 'data_format':nn.data_format, 'resolution': resolution},
{'types' : (t.IMG_TRANSFORMED, face_type, t.MODE_M), 'data_format':nn.data_format, 'resolution': resolution} ],
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)
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)
return ( ('src_loss', src_loss), ('dst_loss', dst_loss), )
@ -728,10 +718,12 @@ class SAEHDModel(ModelBase):
for sample_list in samples ]
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] ]
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 = []
st = []
@ -753,12 +745,16 @@ class SAEHDModel(ModelBase):
return result
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']:
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]
return bgr[0], mask[...,0]
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]
#override

2
requirements-colab.txt
View file

@ -6,4 +6,4 @@ ffmpeg-python==0.1.17
scikit-image==0.14.2
scipy==1.4.1
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
scipy==1.4.1
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:
samples_cache = dict()
@staticmethod
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')
normalize_tanh = opts.get('normalize_tanh', False)
data_format = opts.get('data_format', 'NHWC')
img_type = SPTF.NONE
target_face_type = SPTF.NONE
@ -300,6 +302,10 @@ class SampleProcessor(object):
else:
img = np.clip (img, 0.0, 1.0)
if data_format == "NCHW":
img = np.transpose(img, (2,0,1) )
outputs_sample.append ( img )
outputs += [outputs_sample]