diff --git a/joblib/SubprocessorBase.py b/joblib/SubprocessorBase.py index d2da803..bf8cd86 100644 --- a/joblib/SubprocessorBase.py +++ b/joblib/SubprocessorBase.py @@ -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") diff --git a/models/Model_SAE/Model.py b/models/Model_SAE/Model.py index 644ea0c..2616461 100644 --- a/models/Model_SAE/Model.py +++ b/models/Model_SAE/Model.py @@ -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 diff --git a/nnlib/CAInitializer.py b/nnlib/CAInitializer.py new file mode 100644 index 0000000..20c03c1 --- /dev/null +++ b/nnlib/CAInitializer.py @@ -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) \ No newline at end of file diff --git a/nnlib/nnlib.py b/nnlib/nnlib.py index 315387e..dbf0544 100644 --- a/nnlib/nnlib.py +++ b/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 + \ No newline at end of file