diff --git a/models/ModelBase.py b/models/ModelBase.py index 173e33e..637d761 100644 --- a/models/ModelBase.py +++ b/models/ModelBase.py @@ -121,7 +121,8 @@ class ModelBase(object): nnlib.import_all ( nnlib.DeviceConfig(allow_growth=False, **self.device_args) ) self.device_config = nnlib.active_DeviceConfig self.keras = nnlib.keras - + self.K = nnlib.keras.backend + self.onInitialize() self.options['batch_size'] = self.batch_size @@ -282,17 +283,20 @@ class ModelBase(object): if len(optimizer_filename_list) != 0: opt_filename = self.get_strpath_storage_for_file('opt.h5') if Path(opt_filename).exists(): - h5dict = self.keras.utils.io_utils.H5Dict(opt_filename, mode='r') try: + with open(opt_filename, "rb") as f: + d = pickle.loads(f.read()) + for x in optimizer_filename_list: opt, filename = x - if filename in h5dict: - opt = opt.__class__.from_config( json.loads(h5dict[filename]) ) - x[0] = opt - finally: - h5dict.close() - - return [x[0] for x in optimizer_filename_list] + if filename in d: + weights = d[filename].get('weights', None) + if weights: + opt.set_weights(weights) + print("set ok") + except Exception as e: + print ("Unable to load ", opt_filename) + def save_weights_safe(self, model_filename_list, optimizer_filename_list=[]): for model, filename in model_filename_list: @@ -302,13 +306,23 @@ class ModelBase(object): rename_list = model_filename_list if len(optimizer_filename_list) != 0: opt_filename = self.get_strpath_storage_for_file('opt.h5') - h5dict = self.keras.utils.io_utils.H5Dict(opt_filename + '.tmp', mode='w') + try: + d = {} for opt, filename in optimizer_filename_list: - h5dict[filename] = json.dumps(opt.get_config()) - finally: - h5dict.close() + fd = {} + symbolic_weights = getattr(opt, 'weights') + if symbolic_weights: + fd['weights'] = self.K.batch_get_value(symbolic_weights) + + d[filename] = fd + + with open(opt_filename+'.tmp', 'wb') as f: + f.write( pickle.dumps(d) ) + rename_list += [('', 'opt.h5')] + except Exception as e: + print ("Unable to save ", opt_filename) for _, filename in rename_list: filename = self.get_strpath_storage_for_file(filename) diff --git a/models/Model_SAE/Model.py b/models/Model_SAE/Model.py index 965fd02..8d05743 100644 --- a/models/Model_SAE/Model.py +++ b/models/Model_SAE/Model.py @@ -223,13 +223,6 @@ class SAEModel(ModelBase): pred_dst_dstm = self.decoder_dstm(warped_dst_code) pred_src_dstm = self.decoder_srcm(warped_dst_code) - self.src_dst_opt, \ - self.src_dst_mask_opt = self.load_weights_safe( - weights_to_load, - [ [Adam(lr=5e-5, beta_1=0.5, beta_2=0.999), 'src_dst_opt'], - [Adam(lr=5e-5, beta_1=0.5, beta_2=0.999), 'src_dst_mask_opt'] - ]) - 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']: @@ -267,6 +260,9 @@ class SAEModel(ModelBase): psd_target_dst_anti_masked_ar = [ pred_src_dst_sigm_ar[i]*target_dstm_anti_sigm_ar[i] for i in range(len(pred_src_dst_sigm_ar))] if self.is_training_mode: + self.src_dst_opt = Adam(lr=5e-5, beta_1=0.5, beta_2=0.999) + self.src_dst_mask_opt = Adam(lr=5e-5, beta_1=0.5, beta_2=0.999) + if self.options['archi'] == 'liae': src_dst_loss_train_weights = self.encoder.trainable_weights + self.inter_B.trainable_weights + self.inter_AB.trainable_weights + self.decoder.trainable_weights if self.options['learn_mask']: @@ -325,14 +321,17 @@ class SAEModel(ModelBase): self.AE_view = K.function ([warped_src, warped_dst], [pred_src_src[-1], pred_dst_dst[-1], pred_src_dst[-1], pred_src_dstm[-1]]) else: self.AE_view = K.function ([warped_src, warped_dst], [pred_src_src[-1], pred_dst_dst[-1], pred_src_dst[-1] ] ) - + + self.load_weights_safe(weights_to_load)#, [ [self.src_dst_opt, 'src_dst_opt'], [self.src_dst_mask_opt, 'src_dst_mask_opt']]) else: + self.load_weights_safe(weights_to_load) if self.options['learn_mask']: self.AE_convert = K.function ([warped_dst],[ pred_src_dst[-1], pred_src_dstm[-1] ]) else: self.AE_convert = K.function ([warped_dst],[ pred_src_dst[-1] ]) - - if self.is_training_mode: + + + if self.is_training_mode: self.src_sample_losses = [] self.dst_sample_losses = [] @@ -353,6 +352,7 @@ class SAEModel(ModelBase): sample_process_options=SampleProcessor.Options(random_flip=self.random_flip, normalize_tanh = True), output_sample_types=output_sample_types ) ]) + #override def onSave(self): opt_ar = [ [self.src_dst_opt, 'src_dst_opt'], diff --git a/nnlib/nnlib.py b/nnlib/nnlib.py index db0cb7c..0debf31 100644 --- a/nnlib/nnlib.py +++ b/nnlib/nnlib.py @@ -71,8 +71,8 @@ ZeroPadding2D = keras.layers.ZeroPadding2D RandomNormal = keras.initializers.RandomNormal Model = keras.models.Model -#Adam = keras.optimizers.Adam -Adam = nnlib.Adam +Adam = keras.optimizers.Adam +FastAdam = nnlib.FastAdam modelify = nnlib.modelify gaussian_blur = nnlib.gaussian_blur @@ -434,21 +434,16 @@ NLayerDiscriminator = nnlib.NLayerDiscriminator return dict(list(base_config.items()) + list(config.items())) nnlib.Scale = Scale - class Adam(keras.optimizers.Optimizer): - def __init__(self, lr=0.001, beta_1=0.9, beta_2=0.999, - epsilon=None, decay=0., amsgrad=False, iterations=0, **kwargs): - super(Adam, self).__init__(**kwargs) + class FastAdam(keras.optimizers.Optimizer): + def __init__(self, lr=0.001, beta_1=0.9, beta_2=0.999, iterations=0, **kwargs): + super(FastAdam, self).__init__(**kwargs) with K.name_scope(self.__class__.__name__): self.iterations = K.variable(iterations, 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.initial_decay = decay - self.amsgrad = amsgrad + + self.epsilon = K.epsilon() @keras.legacy.interfaces.legacy_get_updates_support def get_updates(self, loss, params): @@ -456,34 +451,16 @@ NLayerDiscriminator = nnlib.NLayerDiscriminator 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))) + (1. - K.pow(self.beta_1, t))) + self.weights = [self.iterations] - 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 + for p, g in zip(params, grads): - for p, g, m, v, vhat in zip(params, grads, ms, vs, vhats): - 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) - p_t = p - lr_t * m_t / (K.sqrt(vhat_t) + self.epsilon) - self.updates.append(K.update(vhat, vhat_t)) - else: - p_t = p - lr_t * m_t / (K.sqrt(v_t) + self.epsilon) - - self.updates.append(K.update(m, m_t)) - self.updates.append(K.update(v, v_t)) + m_t = (1. - self.beta_1) * g + v_t = (1. - self.beta_2) * K.square(g) + p_t = p - lr_t * m_t / (K.sqrt(v_t) + self.epsilon) new_p = p_t # Apply constraints. @@ -497,15 +474,14 @@ NLayerDiscriminator = nnlib.NLayerDiscriminator config = {'iterations': int(K.get_value(self.iterations)), '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} - base_config = super(Adam, self).get_config() + 'beta_2': float(K.get_value(self.beta_2)) + } + base_config = super(FastAdam, self).get_config() return dict(list(base_config.items()) + list(config.items())) - nnlib.Adam = Adam - ''' + nnlib.FastAdam = FastAdam + + ''' not implemented in plaidML class ReflectionPadding2D(keras.layers.Layer): def __init__(self, padding=(1, 1), **kwargs):