mnist_eager.py 7.36 KB
Newer Older
Asim Shankar's avatar
Asim Shankar committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
# Copyright 2018 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.
# ==============================================================================
"""MNIST model training with TensorFlow eager execution.

See:
https://research.googleblog.com/2017/10/eager-execution-imperative-define-by.html

This program demonstrates training of the convolutional neural network model
defined in mnist.py with eager execution enabled.

If you are not interested in eager execution, you should ignore this file.
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os
import sys
import time

33
34
35
36
37
38
# pylint: disable=g-bad-import-order
from absl import app as absl_app
from absl import flags
import tensorflow as tf
import tensorflow.contrib.eager as tfe
# pylint: enable=g-bad-import-order
39

Karmel Allison's avatar
Karmel Allison committed
40
from official.mnist import dataset as mnist_dataset
41
from official.mnist import mnist
42
from official.utils.flags import core as flags_core
Asim Shankar's avatar
Asim Shankar committed
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58


def loss(logits, labels):
  return tf.reduce_mean(
      tf.nn.sparse_softmax_cross_entropy_with_logits(
          logits=logits, labels=labels))


def compute_accuracy(logits, labels):
  predictions = tf.argmax(logits, axis=1, output_type=tf.int64)
  labels = tf.cast(labels, tf.int64)
  batch_size = int(logits.shape[0])
  return tf.reduce_sum(
      tf.cast(tf.equal(predictions, labels), dtype=tf.float32)) / batch_size


59
def train(model, optimizer, dataset, step_counter, log_interval=None):
Asim Shankar's avatar
Asim Shankar committed
60
61
62
63
  """Trains model on `dataset` using `optimizer`."""

  start = time.time()
  for (batch, (images, labels)) in enumerate(tfe.Iterator(dataset)):
64
65
    with tf.contrib.summary.record_summaries_every_n_global_steps(
        10, global_step=step_counter):
Asim Shankar's avatar
Asim Shankar committed
66
67
68
      # Record the operations used to compute the loss given the input,
      # so that the gradient of the loss with respect to the variables
      # can be computed.
69
      with tf.GradientTape() as tape:
Asim Shankar's avatar
Asim Shankar committed
70
71
72
73
74
75
        logits = model(images, training=True)
        loss_value = loss(logits, labels)
        tf.contrib.summary.scalar('loss', loss_value)
        tf.contrib.summary.scalar('accuracy', compute_accuracy(logits, labels))
      grads = tape.gradient(loss_value, model.variables)
      optimizer.apply_gradients(
76
          zip(grads, model.variables), global_step=step_counter)
Asim Shankar's avatar
Asim Shankar committed
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
      if log_interval and batch % log_interval == 0:
        rate = log_interval / (time.time() - start)
        print('Step #%d\tLoss: %.6f (%d steps/sec)' % (batch, loss_value, rate))
        start = time.time()


def test(model, dataset):
  """Perform an evaluation of `model` on the examples from `dataset`."""
  avg_loss = tfe.metrics.Mean('loss')
  accuracy = tfe.metrics.Accuracy('accuracy')

  for (images, labels) in tfe.Iterator(dataset):
    logits = model(images, training=False)
    avg_loss(loss(logits, labels))
    accuracy(
        tf.argmax(logits, axis=1, output_type=tf.int64),
        tf.cast(labels, tf.int64))
  print('Test set: Average loss: %.4f, Accuracy: %4f%%\n' %
        (avg_loss.result(), 100 * accuracy.result()))
  with tf.contrib.summary.always_record_summaries():
    tf.contrib.summary.scalar('loss', avg_loss.result())
    tf.contrib.summary.scalar('accuracy', accuracy.result())


101
def main(flags_obj):
102
  tf.enable_eager_execution()
Asim Shankar's avatar
Asim Shankar committed
103

104
  # Automatically determine device and data_format
Asim Shankar's avatar
Asim Shankar committed
105
  (device, data_format) = ('/gpu:0', 'channels_first')
106
  if flags_obj.no_gpu or tf.test.is_gpu_available():
Asim Shankar's avatar
Asim Shankar committed
107
    (device, data_format) = ('/cpu:0', 'channels_last')
108
  # If data_format is defined in FLAGS, overwrite automatically set value.
109
110
  if flags_obj.data_format is not None:
    data_format = flags_obj.data_format
Asim Shankar's avatar
Asim Shankar committed
111
112
113
  print('Using device %s, and data format %s.' % (device, data_format))

  # Load the datasets
114
115
116
117
  train_ds = mnist_dataset.train(flags_obj.data_dir).shuffle(60000).batch(
      flags_obj.batch_size)
  test_ds = mnist_dataset.test(flags_obj.data_dir).batch(
      flags_obj.batch_size)
Asim Shankar's avatar
Asim Shankar committed
118
119

  # Create the model and optimizer
120
  model = mnist.create_model(data_format)
121
  optimizer = tf.train.MomentumOptimizer(flags_obj.lr, flags_obj.momentum)
Asim Shankar's avatar
Asim Shankar committed
122

123
  # Create file writers for writing TensorBoard summaries.
124
  if flags_obj.output_dir:
Asim Shankar's avatar
Asim Shankar committed
125
126
127
    # Create directories to which summaries will be written
    # tensorboard --logdir=<output_dir>
    # can then be used to see the recorded summaries.
128
129
130
    train_dir = os.path.join(flags_obj.output_dir, 'train')
    test_dir = os.path.join(flags_obj.output_dir, 'eval')
    tf.gfile.MakeDirs(flags_obj.output_dir)
Asim Shankar's avatar
Asim Shankar committed
131
132
133
134
135
136
137
  else:
    train_dir = None
    test_dir = None
  summary_writer = tf.contrib.summary.create_file_writer(
      train_dir, flush_millis=10000)
  test_summary_writer = tf.contrib.summary.create_file_writer(
      test_dir, flush_millis=10000, name='test')
138
139

  # Create and restore checkpoint (if one exists on the path)
140
  checkpoint_prefix = os.path.join(flags_obj.model_dir, 'ckpt')
141
142
143
144
  step_counter = tf.train.get_or_create_global_step()
  checkpoint = tfe.Checkpoint(
      model=model, optimizer=optimizer, step_counter=step_counter)
  # Restore variables on creation if a checkpoint exists.
145
  checkpoint.restore(tf.train.latest_checkpoint(flags_obj.model_dir))
146
147

  # Train and evaluate for a set number of epochs.
Asim Shankar's avatar
Asim Shankar committed
148
  with tf.device(device):
149
    for _ in range(flags_obj.train_epochs):
150
151
      start = time.time()
      with summary_writer.as_default():
152
153
        train(model, optimizer, train_ds, step_counter,
              flags_obj.log_interval)
154
155
156
157
158
      end = time.time()
      print('\nTrain time for epoch #%d (%d total steps): %f' %
            (checkpoint.save_counter.numpy() + 1,
             step_counter.numpy(),
             end - start))
Asim Shankar's avatar
Asim Shankar committed
159
160
      with test_summary_writer.as_default():
        test(model, test_ds)
161
      checkpoint.save(checkpoint_prefix)
Asim Shankar's avatar
Asim Shankar committed
162
163


164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
def define_mnist_eager_flags():
  """Defined flags and defaults for MNIST in eager mode."""
  flags_core.define_base_eager()
  flags_core.define_image()
  flags.adopt_module_key_flags(flags_core)

  flags.DEFINE_integer(
      name='log_interval', short_name='li', default=10,
      help=flags_core.help_wrap('batches between logging training status'))

  flags.DEFINE_string(
      name='output_dir', short_name='od', default=None,
      help=flags_core.help_wrap('Directory to write TensorBoard summaries'))

  flags.DEFINE_float(name='learning_rate', short_name='lr', default=0.01,
                     help=flags_core.help_wrap('Learning rate.'))

  flags.DEFINE_float(name='momentum', short_name='m', default=0.5,
                     help=flags_core.help_wrap('SGD momentum.'))

  flags.DEFINE_bool(name='no_gpu', short_name='nogpu', default=False,
                    help=flags_core.help_wrap(
                        'disables GPU usage even if a GPU is available'))

  flags_core.set_defaults(
      data_dir='/tmp/tensorflow/mnist/input_data',
      model_dir='/tmp/tensorflow/mnist/checkpoints/',
      batch_size=100,
      train_epochs=10,
  )
Asim Shankar's avatar
Asim Shankar committed
194

195
if __name__ == '__main__':
196
197
  define_mnist_eager_flags()
  absl_app.run(main=main)