fruits_net_train.py

# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

# NOTE: Based off of the Advanced CNN tutorial in Tensorflow documentation

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from datetime import datetime
import time

import tensorflow as tf

import fruits_net

FLAGS = tf.flags.FLAGS

tf.flags.DEFINE_string('train_dir', './tmp/fruits_train',
                           """Directory where to write event logs """
                           """and checkpoint.""")
tf.flags.DEFINE_integer('max_steps', 1,
                            """Number of batches to run.""")
tf.flags.DEFINE_boolean('log_device_placement', False,
                            """Whether to log device placement.""")
tf.flags.DEFINE_integer('log_frequency', 1,
                            """How often to log results to the console.""")


def train():
  """Train Fruits-360 for a number of steps."""
  with tf.Graph().as_default():
    global_step = tf.train.get_or_create_global_step()

    # Get images and labels for Fruits-360.
    # Force input pipeline to CPU:0 to avoid operations sometimes ending up on
    # GPU and resulting in a slow down.
    with tf.device('/cpu:0'):
      images, labels = fruits_net.inputs()
    print(tf.Session().run(tf.gather(labels, [i for i in range(FLAGS.batch_size)])))

    # Build a Graph that computes the logits predictions from the
    # inference model.
    logits = fruits_net.inference(images)

    # Calculate loss.
    loss = fruits_net.loss(logits, labels)

    # Build a Graph that trains the model with one batch of examples and
    # updates the model parameters.
    train_op = fruits_net.train(loss, global_step)

    class _LoggerHook(tf.train.SessionRunHook):
      """Logs loss and runtime."""

      def begin(self):
        self._step = -1
        self._start_time = time.time()

      def before_run(self, run_context):
        self._step += 1
        return tf.train.SessionRunArgs(loss)  # Asks for loss value.

      def after_run(self, run_context, run_values):
        if self._step % FLAGS.log_frequency == 0:
          current_time = time.time()
          duration = current_time - self._start_time
          self._start_time = current_time

          loss_value = run_values.results
          examples_per_sec = FLAGS.log_frequency * FLAGS.batch_size / duration
          sec_per_batch = float(duration / FLAGS.log_frequency)

          format_str = ('%s: step %d, loss = %.2f (%.1f examples/sec; %.3f '
                        'sec/batch)')
          print (format_str % (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(num_steps=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 main(argv=None):  # pylint: disable=unused-argument
  if tf.gfile.Exists(FLAGS.train_dir):
    tf.gfile.DeleteRecursively(FLAGS.train_dir)
  tf.gfile.MakeDirs(FLAGS.train_dir)
  num_iterations = int(37836 / FLAGS.batch_size)
  fruits_net.set_sampling_chance(0.1)
  for i in range(num_iterations + 1):
    print("===Static Run %d/%d===" % (i, num_iterations + 1))
    train()
  fruits_net.set_sampling_chance(0.0001)
  for i in range(2 * num_iterations):
    print("===Randomized Run %d/%d===" % (i, 2 * num_iterations))
    train()



if __name__ == '__main__':
  tf.app.run()