Append model shape to the checkpoint directory name

igv · Aug 21, 2017 · f8af70e · f8af70e
1 parent 15bc798
commit f8af70e
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/model.py b/model.py
@@ -226,7 +226,8 @@ def prelu(self, _x, i):
 
   def save(self, checkpoint_dir, step):
     model_name = "FSRCNN.model"
-    model_dir = "%s_%s" % ("fsrcnn", self.label_size)
+    d, s, m = self.model_params
+    model_dir = "%s_%s_%s-%s-%s" % ("fsrcnn", self.label_size, d, s, m)
     checkpoint_dir = os.path.join(checkpoint_dir, model_dir)
 
     if not os.path.exists(checkpoint_dir):
@@ -238,7 +239,8 @@ def save(self, checkpoint_dir, step):
 
   def load(self, checkpoint_dir):
     print(" [*] Reading checkpoints...")
-    model_dir = "%s_%s" % ("fsrcnn", self.label_size)
+    d, s, m = self.model_params
+    model_dir = "%s_%s_%s-%s-%s" % ("fsrcnn", self.label_size, d, s, m)
     checkpoint_dir = os.path.join(checkpoint_dir, model_dir)
 
     ckpt = tf.train.get_checkpoint_state(checkpoint_dir)