TensorFlow 1.3 "canned estimators" demo

README.md

1. mnist_cnn.py - baseline code (till TF 1.2)
2. mnist_estimator.py - using TF 1.3 tf.estimator API

mnist_cnn.py

#
#   mnist_cnn.py
#       date. 7/15/2017 - refactoring by using tf.layers API
#

import os
import tensorflow as tf
from tensorflow.examples.tutorials.mnist import input_data

mnist = input_data.read_data_sets("../MNIST_data/", one_hot=True)
ckpt_file = '/tmp/mnist_cnn.ckpt'

# Create the model
def cnn_model(x, keep_prob):
    '''
      cnn (convolutional neural network) model
    '''
    net = tf.layers.conv2d(x, 32, kernel_size=(5, 5),
                           strides=(1, 1), padding='same')
    net = tf.layers.max_pooling2d(net, pool_size=(2, 2), 
                                  strides=(2, 2),
                                  padding='same')
    net = tf.layers.conv2d(net, 64, kernel_size=(5, 5),
                           strides=(1, 1), padding='same')
    net = tf.layers.max_pooling2d(net, pool_size=(2, 2), 
                                  strides=(2, 2),
                                  padding='same')
    net = tf.reshape(net, [-1, 7*7*64])
    net = tf.layers.dense(net, 1024, activation=tf.nn.relu)
    net = tf.layers.dropout(net, rate=(1.0 - keep_prob))

    y = tf.layers.dense(net, 10, activation=tf.nn.softmax)

    return y


def inference(x, y_, keep_prob):
    x_image = tf.reshape(x, [-1, 28, 28, 1])
    y_pred = cnn_model(x_image, keep_prob)
    
    loss = tf.reduce_mean(-tf.reduce_sum(y_ * tf.log(y_pred), 
                                    reduction_indices=[1]))
    correct_prediction = tf.equal(tf.argmax(y_pred,1), tf.argmax(y_,1))
    accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))
    
    return loss, accuracy
 
#
if __name__ == '__main__':
    # Variables
    x = tf.placeholder(tf.float32, [None, 784])
    y_ = tf.placeholder(tf.float32, [None, 10])
    keep_prob = tf.placeholder(tf.float32)  
    loss, accuracy = inference(x, y_, keep_prob)

    # Train
    train_step = tf.train.AdagradOptimizer(0.01).minimize(loss)
    vars_to_train = tf.trainable_variables()
    ckpt_index = ckpt_file + '.index'

    if os.path.exists(ckpt_index) == False:
        restore_call = False
        init = tf.global_variables_initializer()
    else:
        restore_call = True
        vars_all = tf.global_variables()
        vars_to_init = list(set(vars_all) - set(vars_to_train))
        init = tf.variables_initializer(vars_to_init)

    saver = tf.train.Saver(vars_to_train)

    with tf.Session() as sess:
        sess.run(init)
        if restore_call:
            # Restore variables from disk.
            saver.restore(sess, ckpt_file) 
            print('model restored.')

        print('Training...')
        for i in range(5001):
            batch_xs, batch_ys = mnist.train.next_batch(100)
            train_step.run({x: batch_xs, y_: batch_ys, keep_prob: 0.5})
            if i % 1000 == 0:
                train_accuracy = accuracy.eval(
                    {x: batch_xs, y_: batch_ys, keep_prob: 1.0})
                print('  step, accurary = %6d: %8.4f' % (i, train_accuracy))

        # Test trained model
        print('accuracy = ', accuracy.eval(
            {x: mnist.test.images, y_: mnist.test.labels, keep_prob: 1.0}))
        # check dataflow status
        print('epoch comp = %4d times' % mnist.train.epochs_completed)
    
        # Save the variables to disk.
        save_path = saver.save(sess, ckpt_file)
        print("Model saved in file: %s" % save_path)

mnist_estimator.py

#
#   mnist_estimator.py
#       mnist demo to use tf.estimator class library
#       (ref.) https://github.com/tensorflow/tensorflow/blob/master/tensorflow/examples/learn/mnist.py
#

import numpy as np
import tensorflow as tf

N_DIGITS = 10       # Number of digits.
X_FEATURE = 'x'     # Name of the input feature.


def cnn_model(features, labels, mode):
    """2-layer convolution model."""
    # Reshape feature to 4d tensor with 2nd and 3rd dimensions being
    # image width and height final dimension being the number of color channels.
    feature = tf.reshape(features[X_FEATURE], [-1, 28, 28, 1])

    # First conv layer will compute 32 features for each 5x5 patch
    with tf.variable_scope('conv_layer1'):
        h_conv1 = tf.layers.conv2d(
            feature,
            filters=32,
            kernel_size=[5, 5],
            padding='same',
            activation=tf.nn.relu)
        h_pool1 = tf.layers.max_pooling2d(
            h_conv1, pool_size=2, strides=2, padding='same')

    # Second conv layer will compute 64 features for each 5x5 patch.
    with tf.variable_scope('conv_layer2'):
        h_conv2 = tf.layers.conv2d(
            h_pool1,
            filters=64,
            kernel_size=[5, 5],
            padding='same',
            activation=tf.nn.relu)
        h_pool2 = tf.layers.max_pooling2d(
            h_conv2, pool_size=2, strides=2, padding='same')
        # reshape tensor into a batch of vectors
        h_pool2_flat = tf.reshape(h_pool2, [-1, 7 * 7 * 64])

    # Densely connected layer with 1024 neurons.
    h_fc1 = tf.layers.dense(h_pool2_flat, 1024, activation=tf.nn.relu)
    if mode == tf.estimator.ModeKeys.TRAIN:
        h_fc1 = tf.layers.dropout(h_fc1, rate=0.5)

    # Compute logits (1 per class) and compute loss.
    logits = tf.layers.dense(h_fc1, N_DIGITS, activation=None)

    return logits


def inference_fn(features, labels, mode):
    # calculate logits
    logits = cnn_model(features, labels, mode)

    # Compute predictions.
    predicted_classes = tf.argmax(logits, 1)
    if mode == tf.estimator.ModeKeys.PREDICT:
        predictions = {
            'class': predicted_classes,
            'prob': tf.nn.softmax(logits)
        }
        return tf.estimator.EstimatorSpec(mode, predictions=predictions)

    # Compute loss.
    onehot_labels = tf.one_hot(tf.cast(labels, tf.int32), N_DIGITS, 1, 0)
    loss = tf.losses.softmax_cross_entropy(
        onehot_labels=onehot_labels, logits=logits)

    # Create training op.
    if mode == tf.estimator.ModeKeys.TRAIN:
        optimizer = tf.train.AdagradOptimizer(learning_rate=0.01)
        train_op = optimizer.minimize(loss, global_step=tf.train.get_global_step())
        ret = tf.estimator.EstimatorSpec(mode, loss=loss, train_op=train_op)
    
    else:   # Compute evaluation metrics.
        eval_metric_ops = {
            'accuracy': tf.metrics.accuracy(
                labels=labels, predictions=predicted_classes)
        }
        ret = tf.estimator.EstimatorSpec(
                mode, loss=loss, eval_metric_ops=eval_metric_ops)
    return ret


def main(unused_args):
    ### Load MNIST dataset.
    mnist = tf.contrib.learn.datasets.DATASETS['mnist']('../MNIST_data')
    train_input_fn = tf.estimator.inputs.numpy_input_fn(
        x={X_FEATURE: mnist.train.images},
        y=mnist.train.labels.astype(np.int32),
        batch_size=100,
        num_epochs=None,
        shuffle=True)
    test_input_fn = tf.estimator.inputs.numpy_input_fn(
        x={X_FEATURE: mnist.train.images},
        y=mnist.train.labels.astype(np.int32),
        num_epochs=1,
        shuffle=False)

    ### Convolutional network
    classifier = tf.estimator.Estimator(model_fn=inference_fn)
    classifier.train(input_fn=train_input_fn, steps=400)
    scores = classifier.evaluate(input_fn=test_input_fn)
    print('Accuracy (conv_model): {0:f}'.format(scores['accuracy']))


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