app bugz

1 files changed, 7 insertions, 8 deletions

diff --git a/app/relay/modules/pix2pixhd.js b/app/relay/modules/pix2pixhd.js
index ccfd334..d5d6ada 100644
--- a/app/relay/modules/pix2pixhd.js
+++ b/app/relay/modules/pix2pixhd.js
@@ -48,14 +48,14 @@ const train = {
     const dataset = task.dataset.toLowerCase()
     const datasets_path = path.join(cwd, 'datasets', dataset)
     const checkpoints_path = path.join(cwd, 'checkpoints', dataset)
+    const iter_txt = path.join(checkpoints_path, 'iter.txt')
+    console.log(dataset, datasets_path, checkpoints_path)
     if (fs.existsSync(checkpoints_path)) {
-      try {
-        const checkpoints = fs.readdirSync(checkpoints_path)
-        checkpoints.forEach(name => {
-          epoch = Math.max(parseInt(name) || 0, epoch)
-        })
-        console.log(task.module, dataset, epoch, task.epochs)
-      } catch (e) { }
+      const iter = fs.readFileSync(iter_txt).split('\n')
+      epoch = iter[0] || 0
+      console.log(task.module, dataset, '=>', epoch, task.epochs)
+    } else {
+      console.log(task.module, dataset, '=>', 'starting new training')
     }
     let args = [
       '--dataroot', datasets_path,
@@ -68,7 +68,6 @@ const train = {
     ]
     if (epoch) {
       args = args.concat([
-        '--epoch_count', task.epoch + task.epochs + 1,
         '--which_epoch', 'latest',
         '--continue_train',
       ])