check if checkpoint exists

video_prediction_tools/main_scripts/main_train_models.py

@@ -259,7 +259,7 @@ class TrainModel(object):
         """
         Restore the train and validation losses in the pickle file if checkpoint is given 
         """
-        if self.start_step == 0:
+        if os.path.exists(os.path.join(self.output_dir,"checkpoint")):
             train_losses = []
             val_losses = []
         else:
@@ -278,6 +278,7 @@ class TrainModel(object):
             print("parameter_count =", sess.run(self.parameter_count))
             sess.run(tf.global_variables_initializer())
             sess.run(tf.local_variables_initializer())
+            if os.path.exists(os.path.join(self.output_dir,"checkpoint")):
                 self.restore(sess, self.checkpoint)
             #sess.graph.finalize()
             self.start_step = sess.run(self.global_step)