def train():
"""Train CIFAR-10 for a number of steps."""
with tf.Graph().as_default():
global_step = contrib_framework.get_or_create_global_step()
# Get images and labels for CIFAR-10.
images, labels = cifar10.distorted_inputs()
# Create a compression object using the compression hyperparameters
compression_obj = cifar10.create_compressor(FLAGS.compression_hparams,
global_step=global_step)
# Build a Graph that computes the logits predictions from the
# inference model.
logits = cifar10.inference(images, compression_obj)
# Calculate loss.
loss = cifar10.loss(logits, labels)
# Build a Graph that trains the model with one batch of examples and
# updates the model parameters.
train_op = cifar10.train(loss, global_step, compression_obj)
class _LoggerHook(tf.train.SessionRunHook):
"""Logs loss and runtime."""
def begin(self):
self._step = -1
def before_run(self, run_context):
self._step += 1
self._start_time = time.time()
return tf.train.SessionRunArgs(loss) # Asks for loss value.
def after_run(self, run_context, run_values):
duration = time.time() - self._start_time
loss_value = run_values.results
if self._step % 10 == 0:
num_examples_per_step = 128
examples_per_sec = num_examples_per_step / duration
sec_per_batch = float(duration)
format_str = (
'%s: step %d, loss = %.2f (%.1f examples/sec; %.3f '
'sec/batch)')
print(format_str %
(datetime.datetime.now(), self._step, loss_value,
examples_per_sec, sec_per_batch))
with tf.train.MonitoredTrainingSession(
checkpoint_dir=FLAGS.train_dir,
hooks=[
tf.train.StopAtStepHook(last_step=FLAGS.max_steps),
tf.train.NanTensorHook(loss),
_LoggerHook()
],
config=tf.ConfigProto(log_device_placement=FLAGS.
log_device_placement)) as mon_sess:
while not mon_sess.should_stop():
mon_sess.run(train_op)
def eval_once():
"""Run Eval once."""
with tf.Graph().as_default() as g:
with tf.Session() as sess:
ckpt = tf.train.get_checkpoint_state(FLAGS.checkpoint_dir)
if not (ckpt and ckpt.model_checkpoint_path):
print('No checkpoint file found')
return
# Get images and labels for CIFAR-10.
eval_data = FLAGS.eval_data == 'test'
images, labels = cifar10.inputs(eval_data=eval_data)
global_step = FLAGS.global_step
# If invalid global step or none provided, use the last global step
# recorded in the checkpoint file.
if not global_step or global_step < 0:
# Assuming model_checkpoint_path looks something like:
# /my-favorite-path/cifar10_train/model.ckpt-0,
# extract global_step from it.
global_step = ckpt.model_checkpoint_path.split('/')[-1].split(
'-')[-1]
compression_obj = cifar10.create_compressor(
FLAGS.compression_hparams, global_step=global_step)
# Build a Graph that computes the logits predictions from the
# inference model.
logits = cifar10.inference(images, compression_obj)
# Calculate predictions.
top_k_op = tf.nn.in_top_k(logits, labels, 1)
# Restore the moving average version of the learned variables for eval.
variable_averages = tf.train.ExponentialMovingAverage(
cifar10.MOVING_AVERAGE_DECAY)
variables_to_restore = variable_averages.variables_to_restore()
saver = tf.train.Saver(variables_to_restore)
# Restores from checkpoint
saver.restore(sess, ckpt.model_checkpoint_path)
# Build the summary operation based on the TF collection of Summaries.
summary_op = tf.summary.merge_all()
summary_writer = tf.summary.FileWriter(FLAGS.eval_dir, g)
# Start the queue runners.
coord = tf.train.Coordinator()
try:
threads = []
for qr in tf.get_collection(tf.GraphKeys.QUEUE_RUNNERS):
threads.extend(
qr.create_threads(sess,
coord=coord,
daemon=True,
start=True))
num_iter = int(math.ceil(FLAGS.num_examples / 128))
true_count = 0 # Counts the number of correct predictions.
total_sample_count = num_iter * 128
step = 0
while step < num_iter and not coord.should_stop():
predictions = sess.run([top_k_op])
true_count += np.sum(predictions)
step += 1
# Compute precision @ 1.
precision = true_count / total_sample_count
print('%s: precision @ 1 = %.3f' %
(datetime.datetime.now(), precision))
summary = tf.Summary()
summary.ParseFromString(sess.run(summary_op))
summary.value.add(tag='Precision @ 1', simple_value=precision)
summary_writer.add_summary(summary, global_step)
except Exception as e: # pylint: disable=broad-except
coord.request_stop(e)
coord.request_stop()
coord.join(threads, stop_grace_period_secs=10)