Fixed initialisation issue erasing loaded model.

base/base_model.py

@@ -22,6 +22,8 @@ def load(self, sess):
             print("Loading model checkpoint {} ...\n".format(latest_checkpoint))
             self.saver.restore(sess, latest_checkpoint)
             print("Model loaded")
+        else:
+            print("NO model loaded. Training from beginning")
 
     # just initialize a tensorflow variable to use it as epoch counter
     def init_cur_epoch(self):

base/base_train.py

@@ -2,7 +2,7 @@
 
 
 class BaseTrain:
-    def __init__(self, sess, model, data, config, logger):
+    def __init__(self, sess, model, data, config, logger, load=False):
         self.model = model
         self.logger = logger
         self.config = config
@@ -11,6 +11,10 @@ def __init__(self, sess, model, data, config, logger):
         self.init = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer())
         self.sess.run(self.init)
 
+        # Load the model after initialization to ensure that the loaded values are kept
+        if load:
+            self.model.load(self.sess)
+
     def train(self):
         for cur_epoch in range(self.model.cur_epoch_tensor.eval(self.sess), self.config.num_epochs + 1, 1):
             self.train_epoch()

mains/example.py

@@ -26,14 +26,12 @@ def main():
     sess = tf.Session()
     # create an instance of the model you want
     model = ExampleModel(config)
-    #load model if exists
-    model.load(sess)
     # create your data generator
     data = DataGenerator(config)
     # create tensorboard logger
     logger = Logger(sess, config)
     # create trainer and pass all the previous components to it
-    trainer = ExampleTrainer(sess, model, data, config, logger)
+    trainer = ExampleTrainer(sess, model, data, config, logger, load=True)
 
     # here you train your model
     trainer.train()

trainers/example_trainer.py

@@ -4,8 +4,8 @@
 
 
 class ExampleTrainer(BaseTrain):
-    def __init__(self, sess, model, data, config,logger):
-        super(ExampleTrainer, self).__init__(sess, model, data, config,logger)
+    def __init__(self, sess, model, data, config,logger, **kwargs):
+        super(ExampleTrainer, self).__init__(sess, model, data, config,logger, **kwargs)
 
     def train_epoch(self):
         loop = tqdm(range(self.config.num_iter_per_epoch))