mirror of
https://github.com/iperov/DeepFaceLab.git
synced 2025-07-07 05:22:06 -07:00
SAE: added option "Use CA weights":
Initialize network with 'Convolution Aware' weights. This may help to achieve a higher accuracy model, but consumes time at first run.
This commit is contained in:
iperov
parent
71ff0ce1a7
commit
d6a45763a2
4 changed files with 346 additions and 18 deletions
|
|
@ -86,7 +86,7 @@ class Subprocessor(object):
|
|||
c2s.put ( {'op': 'error', 'data' : data} )
|
||||
|
||||
#overridable
|
||||
def __init__(self, name, SubprocessorCli_class, no_response_time_sec = 60):
|
||||
def __init__(self, name, SubprocessorCli_class, no_response_time_sec = 0):
|
||||
if not issubclass(SubprocessorCli_class, Subprocessor.Cli):
|
||||
raise ValueError("SubprocessorCli_class must be subclass of Subprocessor.Cli")
|
||||
|
||||
|
|
|
|||
|
|
@ -56,12 +56,15 @@ class SAEModel(ModelBase):
|
|||
|
||||
default_ae_dims = 256 if self.options['archi'] == 'liae' else 512
|
||||
default_ed_ch_dims = 42
|
||||
def_ca_weights = False
|
||||
if is_first_run:
|
||||
self.options['ae_dims'] = np.clip ( io.input_int("AutoEncoder dims (32-1024 ?:help skip:%d) : " % (default_ae_dims) , default_ae_dims, help_message="More dims are better, but requires more VRAM. You can fine-tune model size to fit your GPU." ), 32, 1024 )
|
||||
self.options['ed_ch_dims'] = np.clip ( io.input_int("Encoder/Decoder dims per channel (21-85 ?:help skip:%d) : " % (default_ed_ch_dims) , default_ed_ch_dims, help_message="More dims are better, but requires more VRAM. You can fine-tune model size to fit your GPU." ), 21, 85 )
|
||||
self.options['ca_weights'] = io.input_bool ("Use CA weights? (y/n, ?:help skip: %s ) : " % (yn_str[def_ca_weights]), def_ca_weights, help_message="Initialize network with 'Convolution Aware' weights. This may help to achieve a higher accuracy model, but consumes time at first run.")
|
||||
else:
|
||||
self.options['ae_dims'] = self.options.get('ae_dims', default_ae_dims)
|
||||
self.options['ed_ch_dims'] = self.options.get('ed_ch_dims', default_ed_ch_dims)
|
||||
self.options['ca_weights'] = self.options.get('ca_weights', def_ca_weights)
|
||||
|
||||
if is_first_run:
|
||||
self.options['lighter_encoder'] = io.input_bool ("Use lightweight encoder? (y/n, ?:help skip:n) : ", False, help_message="Lightweight encoder is 35% faster, requires less VRAM, but sacrificing overall quality.")
|
||||
|
|
@ -122,6 +125,7 @@ class SAEModel(ModelBase):
|
|||
target_dstm_ar = [ Input ( ( mask_shape[0] // (2**i) ,)*2 + (mask_shape[-1],) ) for i in range(ms_count-1, -1, -1)]
|
||||
|
||||
|
||||
models_list = []
|
||||
weights_to_load = []
|
||||
if self.options['archi'] == 'liae':
|
||||
self.encoder = modelify(SAEModel.LIAEEncFlow(resolution, self.options['lighter_encoder'], ed_ch_dims=ed_ch_dims) ) (Input(bgr_shape))
|
||||
|
|
@ -135,9 +139,12 @@ class SAEModel(ModelBase):
|
|||
|
||||
self.decoder = modelify(SAEModel.LIAEDecFlow (bgr_shape[2],ed_ch_dims=ed_ch_dims//2, multiscale_count=self.ms_count )) (inter_output_Inputs)
|
||||
|
||||
models_list += [self.encoder, self.inter_B, self.inter_AB, self.decoder]
|
||||
|
||||
if self.options['learn_mask']:
|
||||
self.decoderm = modelify(SAEModel.LIAEDecFlow (mask_shape[2],ed_ch_dims=int(ed_ch_dims/1.5) )) (inter_output_Inputs)
|
||||
|
||||
models_list += [self.decoderm]
|
||||
|
||||
if not self.is_first_run():
|
||||
weights_to_load += [ [self.encoder , 'encoder.h5'],
|
||||
[self.inter_B , 'inter_B.h5'],
|
||||
|
|
@ -146,7 +153,7 @@ class SAEModel(ModelBase):
|
|||
]
|
||||
if self.options['learn_mask']:
|
||||
weights_to_load += [ [self.decoderm, 'decoderm.h5'] ]
|
||||
|
||||
|
||||
warped_src_code = self.encoder (warped_src)
|
||||
warped_src_inter_AB_code = self.inter_AB (warped_src_code)
|
||||
warped_src_inter_code = Concatenate()([warped_src_inter_AB_code,warped_src_inter_AB_code])
|
||||
|
|
@ -175,10 +182,13 @@ class SAEModel(ModelBase):
|
|||
self.decoder_src = modelify(SAEModel.DFDecFlow (bgr_shape[2],ed_ch_dims=ed_ch_dims//2, multiscale_count=self.ms_count )) (dec_Inputs)
|
||||
self.decoder_dst = modelify(SAEModel.DFDecFlow (bgr_shape[2],ed_ch_dims=ed_ch_dims//2, multiscale_count=self.ms_count )) (dec_Inputs)
|
||||
|
||||
models_list += [self.encoder, self.decoder_src, self.decoder_dst]
|
||||
|
||||
if self.options['learn_mask']:
|
||||
self.decoder_srcm = modelify(SAEModel.DFDecFlow (mask_shape[2],ed_ch_dims=int(ed_ch_dims/1.5) )) (dec_Inputs)
|
||||
self.decoder_dstm = modelify(SAEModel.DFDecFlow (mask_shape[2],ed_ch_dims=int(ed_ch_dims/1.5) )) (dec_Inputs)
|
||||
|
||||
models_list += [self.decoder_srcm, self.decoder_dstm]
|
||||
|
||||
if not self.is_first_run():
|
||||
weights_to_load += [ [self.encoder , 'encoder.h5'],
|
||||
[self.decoder_src, 'decoder_src.h5'],
|
||||
|
|
@ -188,7 +198,11 @@ class SAEModel(ModelBase):
|
|||
weights_to_load += [ [self.decoder_srcm, 'decoder_srcm.h5'],
|
||||
[self.decoder_dstm, 'decoder_dstm.h5'],
|
||||
]
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
warped_src_code = self.encoder (warped_src)
|
||||
warped_dst_code = self.encoder (warped_dst)
|
||||
pred_src_src = self.decoder_src(warped_src_code)
|
||||
|
|
@ -208,10 +222,13 @@ class SAEModel(ModelBase):
|
|||
self.decoder_src = modelify(SAEModel.VGDecFlow (bgr_shape[2],ed_ch_dims=ed_ch_dims//2 )) (dec_Inputs)
|
||||
self.decoder_dst = modelify(SAEModel.VGDecFlow (bgr_shape[2],ed_ch_dims=ed_ch_dims//2 )) (dec_Inputs)
|
||||
|
||||
models_list += [self.encoder, self.decoder_src, self.decoder_dst]
|
||||
|
||||
if self.options['learn_mask']:
|
||||
self.decoder_srcm = modelify(SAEModel.VGDecFlow (mask_shape[2],ed_ch_dims=int(ed_ch_dims/1.5) )) (dec_Inputs)
|
||||
self.decoder_dstm = modelify(SAEModel.VGDecFlow (mask_shape[2],ed_ch_dims=int(ed_ch_dims/1.5) )) (dec_Inputs)
|
||||
|
||||
models_list += [self.decoder_srcm, self.decoder_dstm]
|
||||
|
||||
if not self.is_first_run():
|
||||
weights_to_load += [ [self.encoder , 'encoder.h5'],
|
||||
[self.decoder_src, 'decoder_src.h5'],
|
||||
|
|
@ -233,7 +250,17 @@ class SAEModel(ModelBase):
|
|||
pred_src_srcm = self.decoder_srcm(warped_src_code)
|
||||
pred_dst_dstm = self.decoder_dstm(warped_dst_code)
|
||||
pred_src_dstm = self.decoder_srcm(warped_dst_code)
|
||||
|
||||
|
||||
if self.is_first_run() and self.options['ca_weights']:
|
||||
io.log_info ("Initializing CA weights...")
|
||||
conv_weights_list = []
|
||||
for model in models_list:
|
||||
for layer in model.layers:
|
||||
if type(layer) == Conv2D:
|
||||
conv_weights_list += [layer.weights[0]] #Conv2D kernel_weights
|
||||
CAInitializerMP ( conv_weights_list )
|
||||
|
||||
|
||||
pred_src_src, pred_dst_dst, pred_src_dst, = [ [x] if type(x) != list else x for x in [pred_src_src, pred_dst_dst, pred_src_dst, ] ]
|
||||
|
||||
if self.options['learn_mask']:
|
||||
|
|
@ -468,9 +495,6 @@ class SAEModel(ModelBase):
|
|||
def initialize_nn_functions():
|
||||
exec (nnlib.import_all(), locals(), globals())
|
||||
|
||||
def conv_initializer():
|
||||
return RandomNormal(0, 0.02)
|
||||
|
||||
class ResidualBlock(object):
|
||||
def __init__(self, filters, kernel_size=3, padding='same', use_reflection_padding=False):
|
||||
self.filters = filters
|
||||
|
|
@ -484,13 +508,13 @@ class SAEModel(ModelBase):
|
|||
#if self.use_reflection_padding:
|
||||
# #var_x = ReflectionPadding2D(stride=1, kernel_size=kernel_size)(var_x)
|
||||
|
||||
var_x = Conv2D(self.filters, kernel_size=self.kernel_size, padding=self.padding, kernel_initializer=conv_initializer() )(var_x)
|
||||
var_x = Conv2D(self.filters, kernel_size=self.kernel_size, padding=self.padding, kernel_initializer=RandomNormal(0, 0.02) )(var_x)
|
||||
var_x = LeakyReLU(alpha=0.2)(var_x)
|
||||
|
||||
#if self.use_reflection_padding:
|
||||
# #var_x = ReflectionPadding2D(stride=1, kernel_size=kernel_size)(var_x)
|
||||
|
||||
var_x = Conv2D(self.filters, kernel_size=self.kernel_size, padding=self.padding, kernel_initializer=conv_initializer() )(var_x)
|
||||
var_x = Conv2D(self.filters, kernel_size=self.kernel_size, padding=self.padding, kernel_initializer=RandomNormal(0, 0.02) )(var_x)
|
||||
var_x = Scale(gamma_init=keras.initializers.Constant(value=0.1))(var_x)
|
||||
var_x = Add()([var_x, inp])
|
||||
var_x = LeakyReLU(alpha=0.2)(var_x)
|
||||
|
|
@ -499,25 +523,25 @@ class SAEModel(ModelBase):
|
|||
|
||||
def downscale (dim):
|
||||
def func(x):
|
||||
return LeakyReLU(0.1)(Conv2D(dim, kernel_size=5, strides=2, padding='same', kernel_initializer=conv_initializer())(x))
|
||||
return LeakyReLU(0.1)(Conv2D(dim, kernel_size=5, strides=2, padding='same', kernel_initializer=RandomNormal(0, 0.02))(x))
|
||||
return func
|
||||
SAEModel.downscale = downscale
|
||||
|
||||
def downscale_sep (dim):
|
||||
def func(x):
|
||||
return LeakyReLU(0.1)(SeparableConv2D(dim, kernel_size=5, strides=2, padding='same', depthwise_initializer=conv_initializer(), pointwise_initializer=RandomNormal(0, 0.02) )(x))
|
||||
return LeakyReLU(0.1)(SeparableConv2D(dim, kernel_size=5, strides=2, padding='same', depthwise_initializer=RandomNormal(0, 0.02), pointwise_initializer=RandomNormal(0, 0.02) )(x))
|
||||
return func
|
||||
SAEModel.downscale_sep = downscale_sep
|
||||
|
||||
def upscale (dim):
|
||||
def func(x):
|
||||
return SubpixelUpscaler()(LeakyReLU(0.1)(Conv2D(dim * 4, kernel_size=3, strides=1, padding='same', kernel_initializer=conv_initializer() )(x)))
|
||||
return SubpixelUpscaler()(LeakyReLU(0.1)(Conv2D(dim * 4, kernel_size=3, strides=1, padding='same', kernel_initializer=RandomNormal(0, 0.02) )(x)))
|
||||
return func
|
||||
SAEModel.upscale = upscale
|
||||
|
||||
def to_bgr (output_nc):
|
||||
def func(x):
|
||||
return Conv2D(output_nc, kernel_size=5, padding='same', activation='sigmoid', kernel_initializer=conv_initializer() )(x)
|
||||
return Conv2D(output_nc, kernel_size=5, padding='same', activation='sigmoid', kernel_initializer=RandomNormal(0, 0.02) )(x)
|
||||
return func
|
||||
SAEModel.to_bgr = to_bgr
|
||||
|
||||
|
|
|
|||
112
nnlib/CAInitializer.py
Normal file
112
nnlib/CAInitializer.py
Normal file
|
|
@ -0,0 +1,112 @@
|
|||
import numpy as np
|
||||
|
||||
def _compute_fans(shape, data_format='channels_last'):
|
||||
"""Computes the number of input and output units for a weight shape.
|
||||
# Arguments
|
||||
shape: Integer shape tuple.
|
||||
data_format: Image data format to use for convolution kernels.
|
||||
Note that all kernels in Keras are standardized on the
|
||||
`channels_last` ordering (even when inputs are set
|
||||
to `channels_first`).
|
||||
# Returns
|
||||
A tuple of scalars, `(fan_in, fan_out)`.
|
||||
# Raises
|
||||
ValueError: in case of invalid `data_format` argument.
|
||||
"""
|
||||
if len(shape) == 2:
|
||||
fan_in = shape[0]
|
||||
fan_out = shape[1]
|
||||
elif len(shape) in {3, 4, 5}:
|
||||
# Assuming convolution kernels (1D, 2D or 3D).
|
||||
# TH kernel shape: (depth, input_depth, ...)
|
||||
# TF kernel shape: (..., input_depth, depth)
|
||||
if data_format == 'channels_first':
|
||||
receptive_field_size = np.prod(shape[2:])
|
||||
fan_in = shape[1] * receptive_field_size
|
||||
fan_out = shape[0] * receptive_field_size
|
||||
elif data_format == 'channels_last':
|
||||
receptive_field_size = np.prod(shape[:-2])
|
||||
fan_in = shape[-2] * receptive_field_size
|
||||
fan_out = shape[-1] * receptive_field_size
|
||||
else:
|
||||
raise ValueError('Invalid data_format: ' + data_format)
|
||||
else:
|
||||
# No specific assumptions.
|
||||
fan_in = np.sqrt(np.prod(shape))
|
||||
fan_out = np.sqrt(np.prod(shape))
|
||||
return fan_in, fan_out
|
||||
|
||||
def _create_basis(filters, size, floatx, eps_std):
|
||||
if size == 1:
|
||||
return np.random.normal(0.0, eps_std, (filters, size))
|
||||
|
||||
nbb = filters // size + 1
|
||||
li = []
|
||||
for i in range(nbb):
|
||||
a = np.random.normal(0.0, 1.0, (size, size))
|
||||
a = _symmetrize(a)
|
||||
u, _, v = np.linalg.svd(a)
|
||||
li.extend(u.T.tolist())
|
||||
p = np.array(li[:filters], dtype=floatx)
|
||||
return p
|
||||
|
||||
def _symmetrize(a):
|
||||
return a + a.T - np.diag(a.diagonal())
|
||||
|
||||
def _scale_filters(filters, variance):
|
||||
c_var = np.var(filters)
|
||||
p = np.sqrt(variance / c_var)
|
||||
return filters * p
|
||||
|
||||
def CAGenerateWeights ( shape, floatx, data_format, eps_std=0.05, seed=None ):
|
||||
if seed is not None:
|
||||
np.random.seed(seed)
|
||||
|
||||
fan_in, fan_out = _compute_fans(shape, data_format)
|
||||
variance = 2 / fan_in
|
||||
|
||||
rank = len(shape)
|
||||
if rank == 3:
|
||||
row, stack_size, filters_size = shape
|
||||
|
||||
transpose_dimensions = (2, 1, 0)
|
||||
kernel_shape = (row,)
|
||||
correct_ifft = lambda shape, s=[None]: np.fft.irfft(shape, s[0])
|
||||
correct_fft = np.fft.rfft
|
||||
|
||||
elif rank == 4:
|
||||
row, column, stack_size, filters_size = shape
|
||||
|
||||
transpose_dimensions = (2, 3, 1, 0)
|
||||
kernel_shape = (row, column)
|
||||
correct_ifft = np.fft.irfft2
|
||||
correct_fft = np.fft.rfft2
|
||||
|
||||
elif rank == 5:
|
||||
x, y, z, stack_size, filters_size = shape
|
||||
|
||||
transpose_dimensions = (3, 4, 0, 1, 2)
|
||||
kernel_shape = (x, y, z)
|
||||
correct_fft = np.fft.rfftn
|
||||
correct_ifft = np.fft.irfftn
|
||||
else:
|
||||
raise ValueError('rank unsupported')
|
||||
|
||||
kernel_fourier_shape = correct_fft(np.zeros(kernel_shape)).shape
|
||||
|
||||
init = []
|
||||
for i in range(filters_size):
|
||||
basis = _create_basis(
|
||||
stack_size, np.prod(kernel_fourier_shape), floatx, eps_std)
|
||||
basis = basis.reshape((stack_size,) + kernel_fourier_shape)
|
||||
|
||||
filters = [correct_ifft(x, kernel_shape) +
|
||||
np.random.normal(0, eps_std, kernel_shape) for
|
||||
x in basis]
|
||||
|
||||
init.append(filters)
|
||||
|
||||
# Format of array is now: filters, stack, row, column
|
||||
init = np.array(init)
|
||||
init = _scale_filters(init, variance)
|
||||
return init.transpose(transpose_dimensions)
|
||||
196
nnlib/nnlib.py
196
nnlib/nnlib.py
|
|
@ -3,8 +3,13 @@ import sys
|
|||
import contextlib
|
||||
import numpy as np
|
||||
|
||||
from .CAInitializer import CAGenerateWeights
|
||||
import multiprocessing
|
||||
from joblib import Subprocessor
|
||||
|
||||
from utils import std_utils
|
||||
from .device import device
|
||||
from interact import interact as io
|
||||
|
||||
class nnlib(object):
|
||||
device = device #forwards nnlib.devicelib to device in order to use nnlib as standalone lib
|
||||
|
|
@ -73,6 +78,7 @@ Model = keras.models.Model
|
|||
|
||||
#Adam = keras.optimizers.Adam
|
||||
Adam = nnlib.Adam
|
||||
Padam = nnlib.Padam
|
||||
|
||||
modelify = nnlib.modelify
|
||||
gaussian_blur = nnlib.gaussian_blur
|
||||
|
|
@ -82,6 +88,7 @@ dssim = nnlib.dssim
|
|||
PixelShuffler = nnlib.PixelShuffler
|
||||
SubpixelUpscaler = nnlib.SubpixelUpscaler
|
||||
Scale = nnlib.Scale
|
||||
CAInitializerMP = nnlib.CAInitializerMP
|
||||
|
||||
#ReflectionPadding2D = nnlib.ReflectionPadding2D
|
||||
#AddUniformNoise = nnlib.AddUniformNoise
|
||||
|
|
@ -91,7 +98,6 @@ Scale = nnlib.Scale
|
|||
keras_contrib = nnlib.keras_contrib
|
||||
GroupNormalization = keras_contrib.layers.GroupNormalization
|
||||
InstanceNormalization = keras_contrib.layers.InstanceNormalization
|
||||
Padam = keras_contrib.optimizers.Padam
|
||||
"""
|
||||
code_import_dlib_string = \
|
||||
"""
|
||||
|
|
@ -497,7 +503,7 @@ NLayerDiscriminator = nnlib.NLayerDiscriminator
|
|||
|
||||
for p, g, m, v, vhat in zip(params, grads, ms, vs, vhats):
|
||||
e = K.tf.device("/cpu:0") if self.tf_cpu_mode == 2 else None
|
||||
if e: e.__enter__()
|
||||
if e: e.__enter__()
|
||||
m_t = (self.beta_1 * m) + (1. - self.beta_1) * g
|
||||
v_t = (self.beta_2 * v) + (1. - self.beta_2) * K.square(g)
|
||||
|
||||
|
|
@ -532,6 +538,126 @@ NLayerDiscriminator = nnlib.NLayerDiscriminator
|
|||
base_config = super(Adam, self).get_config()
|
||||
return dict(list(base_config.items()) + list(config.items()))
|
||||
nnlib.Adam = Adam
|
||||
|
||||
class Padam(keras.optimizers.Optimizer):
|
||||
"""Partially adaptive momentum estimation optimizer.
|
||||
# Arguments
|
||||
lr: float >= 0. Learning rate.
|
||||
beta_1: float, 0 < beta < 1. Generally close to 1.
|
||||
beta_2: float, 0 < beta < 1. Generally close to 1.
|
||||
epsilon: float >= 0. Fuzz factor. If `None`, defaults to `K.epsilon()`.
|
||||
decay: float >= 0. Learning rate decay over each update.
|
||||
amsgrad: boolean. Whether to apply the AMSGrad variant of this
|
||||
algorithm from the paper "On the Convergence of Adam and
|
||||
Beyond".
|
||||
partial: float, 0 <= partial <= 0.5 . Parameter controlling partial
|
||||
momentum adaption. For `partial=0`, this optimizer behaves like SGD,
|
||||
for `partial=0.5` it behaves like AMSGrad.
|
||||
# References
|
||||
- [Closing the Generalization Gap of Adaptive Gradient Methods
|
||||
in Training Deep Neural Networks](https://arxiv.org/pdf/1806.06763.pdf)
|
||||
"""
|
||||
|
||||
def __init__(self, lr=1e-1, beta_1=0.9, beta_2=0.999,
|
||||
epsilon=1e-8, decay=0., amsgrad=False, partial=1. / 8., tf_cpu_mode=0, **kwargs):
|
||||
if partial < 0 or partial > 0.5:
|
||||
raise ValueError(
|
||||
"Padam: 'partial' must be a positive float with a maximum "
|
||||
"value of `0.5`, since higher values will cause divergence "
|
||||
"during training."
|
||||
)
|
||||
super(Padam, self).__init__(**kwargs)
|
||||
with K.name_scope(self.__class__.__name__):
|
||||
self.iterations = K.variable(0, dtype='int64', name='iterations')
|
||||
self.lr = K.variable(lr, name='lr')
|
||||
self.beta_1 = K.variable(beta_1, name='beta_1')
|
||||
self.beta_2 = K.variable(beta_2, name='beta_2')
|
||||
self.decay = K.variable(decay, name='decay')
|
||||
if epsilon is None:
|
||||
epsilon = K.epsilon()
|
||||
self.epsilon = epsilon
|
||||
self.partial = partial
|
||||
self.initial_decay = decay
|
||||
self.amsgrad = amsgrad
|
||||
self.tf_cpu_mode = tf_cpu_mode
|
||||
|
||||
def get_updates(self, loss, params):
|
||||
grads = self.get_gradients(loss, params)
|
||||
self.updates = [K.update_add(self.iterations, 1)]
|
||||
|
||||
lr = self.lr
|
||||
if self.initial_decay > 0:
|
||||
lr = lr * (1. / (1. + self.decay * K.cast(self.iterations,
|
||||
K.dtype(self.decay))))
|
||||
|
||||
t = K.cast(self.iterations, K.floatx()) + 1
|
||||
lr_t = lr * (K.sqrt(1. - K.pow(self.beta_2, t)) /
|
||||
(1. - K.pow(self.beta_1, t)))
|
||||
e = K.tf.device("/cpu:0") if self.tf_cpu_mode > 0 else None
|
||||
if e: e.__enter__()
|
||||
ms = [K.zeros(K.int_shape(p), dtype=K.dtype(p)) for p in params]
|
||||
vs = [K.zeros(K.int_shape(p), dtype=K.dtype(p)) for p in params]
|
||||
if self.amsgrad:
|
||||
vhats = [K.zeros(K.int_shape(p), dtype=K.dtype(p)) for p in params]
|
||||
else:
|
||||
vhats = [K.zeros(1) for _ in params]
|
||||
self.weights = [self.iterations] + ms + vs + vhats
|
||||
if e: e.__exit__(None, None, None)
|
||||
|
||||
for p, g, m, v, vhat in zip(params, grads, ms, vs, vhats):
|
||||
|
||||
e = K.tf.device("/cpu:0") if self.tf_cpu_mode == 2 else None
|
||||
if e: e.__enter__()
|
||||
|
||||
m_t = (self.beta_1 * m) + (1. - self.beta_1) * g
|
||||
v_t = (self.beta_2 * v) + (1. - self.beta_2) * K.square(g)
|
||||
|
||||
if self.amsgrad:
|
||||
vhat_t = K.maximum(vhat, v_t)
|
||||
self.updates.append(K.update(vhat, vhat_t))
|
||||
|
||||
if e: e.__exit__(None, None, None)
|
||||
|
||||
if self.amsgrad:
|
||||
denom = (K.sqrt(vhat_t) + self.epsilon)
|
||||
else:
|
||||
denom = (K.sqrt(v_t) + self.epsilon)
|
||||
|
||||
self.updates.append(K.update(m, m_t))
|
||||
self.updates.append(K.update(v, v_t))
|
||||
|
||||
# Partial momentum adaption.
|
||||
new_p = p - (lr_t * (m_t / (denom ** (self.partial * 2))))
|
||||
|
||||
# Apply constraints.
|
||||
if getattr(p, 'constraint', None) is not None:
|
||||
new_p = p.constraint(new_p)
|
||||
|
||||
self.updates.append(K.update(p, new_p))
|
||||
return self.updates
|
||||
|
||||
def get_config(self):
|
||||
config = {'lr': float(K.get_value(self.lr)),
|
||||
'beta_1': float(K.get_value(self.beta_1)),
|
||||
'beta_2': float(K.get_value(self.beta_2)),
|
||||
'decay': float(K.get_value(self.decay)),
|
||||
'epsilon': self.epsilon,
|
||||
'amsgrad': self.amsgrad,
|
||||
'partial': self.partial}
|
||||
base_config = super(Padam, self).get_config()
|
||||
return dict(list(base_config.items()) + list(config.items()))
|
||||
nnlib.Padam = Padam
|
||||
|
||||
|
||||
def CAInitializerMP( conv_weights_list ):
|
||||
result = CAInitializerMPSubprocessor ( [ (i, K.int_shape(conv_weights)) for i, conv_weights in enumerate(conv_weights_list) ], K.floatx(), K.image_data_format() ).run()
|
||||
for idx, weights in result:
|
||||
K.set_value ( conv_weights_list[idx], weights )
|
||||
|
||||
nnlib.CAInitializerMP = CAInitializerMP
|
||||
|
||||
|
||||
|
||||
'''
|
||||
not implemented in plaidML
|
||||
class ReflectionPadding2D(keras.layers.Layer):
|
||||
|
|
@ -812,5 +938,71 @@ NLayerDiscriminator = nnlib.NLayerDiscriminator
|
|||
if nnlib.tf is not None:
|
||||
nnlib.tf_sess = None
|
||||
nnlib.tf = None
|
||||
|
||||
|
||||
class CAInitializerMPSubprocessor(Subprocessor):
|
||||
class Cli(Subprocessor.Cli):
|
||||
|
||||
#override
|
||||
def on_initialize(self, client_dict):
|
||||
self.floatx = client_dict['floatx']
|
||||
self.data_format = client_dict['data_format']
|
||||
|
||||
#override
|
||||
def process_data(self, data):
|
||||
idx, shape = data
|
||||
weights = CAGenerateWeights (shape, self.floatx, self.data_format)
|
||||
return idx, weights
|
||||
|
||||
#override
|
||||
def get_data_name (self, data):
|
||||
#return string identificator of your data
|
||||
return "undefined"
|
||||
|
||||
#override
|
||||
def __init__(self, idx_shapes_list, floatx, data_format ):
|
||||
|
||||
self.idx_shapes_list = idx_shapes_list
|
||||
self.floatx = floatx
|
||||
self.data_format = data_format
|
||||
|
||||
self.result = []
|
||||
super().__init__('CAInitializerMP', CAInitializerMPSubprocessor.Cli)
|
||||
|
||||
#override
|
||||
def on_clients_initialized(self):
|
||||
io.progress_bar ("Processing", len (self.idx_shapes_list))
|
||||
|
||||
#override
|
||||
def on_clients_finalized(self):
|
||||
io.progress_bar_close()
|
||||
|
||||
#override
|
||||
def process_info_generator(self):
|
||||
for i in range(multiprocessing.cpu_count()):
|
||||
yield 'CPU%d' % (i), {}, {'device_idx': i,
|
||||
'device_name': 'CPU%d' % (i),
|
||||
'floatx' : self.floatx,
|
||||
'data_format' : self.data_format
|
||||
}
|
||||
|
||||
#override
|
||||
def get_data(self, host_dict):
|
||||
if len (self.idx_shapes_list) > 0:
|
||||
return self.idx_shapes_list.pop(0)
|
||||
|
||||
return None
|
||||
|
||||
#override
|
||||
def on_data_return (self, host_dict, data):
|
||||
self.idx_shapes_list.insert(0, data)
|
||||
|
||||
#override
|
||||
def on_result (self, host_dict, data, result):
|
||||
self.result.append ( result )
|
||||
io.progress_bar_inc(1)
|
||||
|
||||
#override
|
||||
def get_result(self):
|
||||
return self.result
|
||||
|
||||
Loading…
Add table
Add a link
Reference in a new issue