Merge pull request #44 from mozilla/isssue18

Added checkpointing
This commit is contained in:
Kelly Davis 2016-09-28 17:42:39 +02:00 committed by GitHub
commit aeb08b3042

View File

@ -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"
]
},
{
@ -1045,6 +1051,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",
@ -1077,6 +1086,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!\""