revert back Adam

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)        

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'],