From 6b312639fe7efa26e7c6dbdbd1292f4a4909ba63 Mon Sep 17 00:00:00 2001
From: Kelly Davis <kdavis@mozilla.com>
Date: Mon, 26 Sep 2016 17:24:16 +0200
Subject: [PATCH] Added checkpointing

---
 DeepSpeech.ipynb | 19 +++++++++++++++++--
 1 file changed, 17 insertions(+), 2 deletions(-)

diff --git a/DeepSpeech.ipynb b/DeepSpeech.ipynb
index 026d46f4..64493b30 100644
--- a/DeepSpeech.ipynb
+++ b/DeepSpeech.ipynb
@@ -77,6 +77,8 @@
    "outputs": [],
    "source": [
     "import time\n",
+    "import os.path\n",
+    "import tempfile\n",
     "import numpy as np\n",
     "import tensorflow as tf\n",
     "from util.gpu import get_available_gpus\n",
@@ -99,7 +101,9 @@
     "* `learning_rate` - The learning rate we will employ in Adam optimizer[[3]](http://arxiv.org/abs/1412.6980)\n",
     "* `training_iters` - The number of iterations we will train for\n",
     "* `batch_size` - The number of elements in a batch\n",
-    "* `display_step` - The number of iterations we cycle through before displaying progress"
+    "* `display_step` - The number of epochs we cycle through before displaying progress\n",
+    "* `checkpoint_step` - The number of epochs we cycle through before checkpointing the model\n",
+    "* `checkpoint_dir` - The directory in which checkpoints are stored"
    ]
   },
   {
@@ -116,7 +120,9 @@
     "epsilon = 1e-8          # TODO: Determine a reasonable value for this\n",
     "training_iters = 5000   # TODO: Determine a reasonable value for this\n",
     "batch_size = 1          # TODO: Determine a reasonable value for this\n",
-    "display_step = 1        # TODO: Determine a reasonable value for this"
+    "display_step = 1        # TODO: Determine a reasonable value for this\n",
+    "checkpoint_step = 50    # TODO: Determine a reasonable value for this\n",
+    "checkpoint_dir = tempfile.gettempdir() # TODO: Determine a reasonable value for this"
    ]
   },
   {
@@ -981,6 +987,9 @@
     "        # Apply gradients to modify the model\n",
     "        apply_gradient_op = apply_gradients(optimizer, avg_tower_gradients)\n",
     "        \n",
+    "        # Create a saver to checkpoint the model\n",
+    "        saver = tf.train.Saver(tf.all_variables())\n",
+    "        \n",
     "        # Create session in which to execute\n",
     "        session = tf.Session(config=tf.ConfigProto(allow_soft_placement=True, log_device_placement=True))\n",
     "        \n",
@@ -1003,6 +1012,12 @@
     "            # Print progress message\n",
     "            if epoch % display_step == 0:\n",
     "                print \"Epoch:\", '%04d' % (epoch+1), \"avg_cer=\", \"{:.9f}\".format((total_accuracy / total_batch))\n",
+    "            \n",
+    "            # Checkpoint the model\n",
+    "            if (epoch % checkpoint_step == 0) or (epoch == training_iters - 1):\n",
+    "                checkpoint_path = os.path.join(checkpoint_dir, 'model.ckpt')\n",
+    "                print \"Checkpointing in directory\", \"%s\" % checkpoint_dir\n",
+    "                saver.save(session, checkpoint_path, global_step=epoch)\n",
     "        \n",
     "        # Indicate optimization has concluded\n",
     "        print \"Optimization Finished!\""