Added ability to save and load weights via args

Author: ncollier

Pilot1/Uno/uno.py

@@ -191,7 +191,15 @@ def set_locals(self):
     {'name': 'growth_bins',
         'type': int,
         'default': 0,
-        'help': 'number of bins to use when discretizing growth response'}
+        'help': 'number of bins to use when discretizing growth response'},
+    {'name' : 'initial_weights',
+        'type' : str,
+        'default': None,
+        'help' : 'file name of initial weights'},
+    {'name' : 'save_weights',
+        'type': str,
+        'default' : None,
+        'help': 'name of file to save weights to' }
 ]
 
 required = [

Pilot1/Uno/uno_baseline_keras2.py

@@ -201,6 +201,14 @@ def build_feature_model(input_shape, name='', dense_layers=[1000, 1000],
     model = Model(x_input, h, name=name)
     return model
 
+class SimpleWeightSaver(Callback):
+    def __init__(self, fname):
+        self.fname = fname
+
+    def on_train_end(self, logs={}):
+        self.model.save_weights(self.fname)
+    
+
 
 def build_model(loader, args, permanent_dropout=True, silent=False):
     input_models = {}
@@ -386,12 +394,17 @@ def warmup_scheduler(epoch):
             logger.info('Cross validation fold {}/{}:'.format(fold + 1, cv))
             cv_ext = '.cv{}'.format(fold + 1)
 
+        template_model = build_model(loader, args, silent=True)
+        if args.initial_weights:
+            logger.info("Loading weights from {}".format(args.initial_weights))
+            template_model.load_weights(args.initial_weights)
+
         if len(args.gpus) > 1:
             from keras.utils import multi_gpu_model
             gpu_count = len(args.gpus)
-            model = multi_gpu_model(build_model(loader, args, silent=True), cpu_merge=False, gpus=gpu_count)
+            model = multi_gpu_model(template_model, cpu_merge=False, gpus=gpu_count)
         else:
-            model = build_model(loader, args, silent=True)
+            model = template_model
 
         optimizer = optimizers.deserialize({'class_name': args.optimizer, 'config': {}})
         base_lr = args.base_lr or K.get_value(optimizer.lr)
@@ -411,7 +424,7 @@ def warmup_scheduler(epoch):
         checkpointer = MultiGPUCheckpoint(prefix + cv_ext + '.model.h5', save_best_only=True)
         tensorboard = TensorBoard(log_dir="tb/{}{}{}".format(args.tb_prefix, ext, cv_ext))
         history_logger = LoggingCallback(logger.debug)
-
+            
         callbacks = [candle_monitor, timeout_monitor, history_logger]
         if args.reduce_lr:
             callbacks.append(reduce_lr)
@@ -421,6 +434,8 @@ def warmup_scheduler(epoch):
             callbacks.append(checkpointer)
         if args.tb:
             callbacks.append(tensorboard)
+        if args.save_weights:
+            callbacks.append(SimpleWeightSaver(args.save_path + '/' + args.save_weights))
 
         if args.use_exported_data is not None:
             train_gen = DataFeeder(filename=args.use_exported_data, batch_size=args.batch_size, shuffle=args.shuffle, single=args.single, agg_dose=args.agg_dose)