Tensorflow Dataset API initialiser hook fix

mnist_new_dataset_hook.py

from __future__ import division, print_function, absolute_import, \
    unicode_literals

import tensorflow as tf
from tensorflow.examples.tutorials.mnist import input_data as mnist_data
from tensorflow.contrib import slim
from tensorflow.contrib.learn import ModeKeys
from tensorflow.contrib.learn.python.learn import learn_runner


tf.logging.set_verbosity(tf.logging.DEBUG)


# Define data loaders #####################################
MNIST_DATA_FOLDER = './MNIST_data'
mnist = mnist_data.read_data_sets(MNIST_DATA_FOLDER, one_hot=False)


class IteratorInitializerHook(tf.train.SessionRunHook):
    def __init__(self):
        super(IteratorInitializerHook, self).__init__()
        self.iterator_initiliser_func = None

    def after_create_session(self, session, coord):
        self.iterator_initiliser_func(session)


# Define the training inputs
def get_train_inputs(batch_size):
    iterator_initiliser_hook = IteratorInitializerHook()
    def train_inputs():
        with tf.name_scope('Training_data'):
            images = mnist.train.images.reshape([-1, 28, 28, 1])
            labels = mnist.train.labels
            images_placeholder = tf.placeholder(
                images.dtype, images.shape)
            labels_placeholder = tf.placeholder(
                labels.dtype, labels.shape)
            dataset = tf.contrib.data.Dataset.from_tensor_slices(
                (images_placeholder, labels_placeholder))
            dataset = dataset.repeat(None)  # Infinite
            dataset = dataset.shuffle(buffer_size=10000)
            dataset = dataset.batch(batch_size)
            iterator = dataset.make_initializable_iterator()
            next_example, next_label = iterator.get_next()
            # Set runhook to initilise iterator
            iterator_initiliser_hook.iterator_initiliser_func = \
                lambda sess: sess.run(
                    iterator.initializer,
                    feed_dict={images_placeholder: images,
                               labels_placeholder: labels})
            return next_example, next_label
    return train_inputs, iterator_initiliser_hook



def get_test_inputs(batch_size):
    iterator_initiliser_hook = IteratorInitializerHook()
    def test_inputs():
        with tf.name_scope('Test_data'):
            images = mnist.test.images.reshape([-1, 28, 28, 1])
            labels = mnist.test.labels
            images_placeholder = tf.placeholder(
                images.dtype, images.shape)
            labels_placeholder = tf.placeholder(
                labels.dtype, labels.shape)
            dataset = tf.contrib.data.Dataset.from_tensor_slices(
                (images_placeholder, labels_placeholder))
            dataset = dataset.batch(batch_size)
            iterator = dataset.make_initializable_iterator()
            next_example, next_label = iterator.get_next()
            # Set runhook to initilise iterator
            iterator_initiliser_hook.iterator_initiliser_func = \
                lambda sess: sess.run(
                    iterator.initializer,
                    feed_dict={images_placeholder: images,
                               labels_placeholder: labels})
            return next_example, next_label
    return test_inputs, iterator_initiliser_hook


# Define model ############################################
def get_estimator(run_config, hparams):
    return tf.contrib.learn.Estimator(
        model_fn=get_model_fn,
        params=hparams,
        config=run_config
    )


def get_model_fn(inputs, targets, mode, params):
    # Define model's architecture
    logits = architecture(inputs, is_training=mode == ModeKeys.TRAIN)
    head = tf.contrib.learn.multi_class_head(
        n_classes=params.n_classes,
        loss_fn=tf.losses.sparse_softmax_cross_entropy)
    return head.create_model_fn_ops(
        features={'inputs': inputs},
        labels=tf.cast(targets, tf.int32),
        mode=mode,
        train_op_fn=get_train_op_fn(params),
        logits=logits
    )


def get_train_op_fn(params):
    def train_op_fn(loss):
        return tf.contrib.layers.optimize_loss(
            loss=loss,
            global_step=tf.contrib.framework.get_global_step(),
            optimizer=tf.train.AdamOptimizer,
            learning_rate=params.learning_rate
        )
    return train_op_fn


def architecture(inputs, is_training, scope='MnistConvNet'):
    tf.logging.debug('is_training: {}, {}'.format(type(is_training), is_training))
    with tf.variable_scope(scope):
        with slim.arg_scope(
                [slim.conv2d, slim.fully_connected],
                weights_initializer=tf.contrib.layers.xavier_initializer()):
            net = slim.conv2d(inputs, 20, [5, 5], padding='VALID',
                              scope='layer1-conv')
            net = slim.max_pool2d(net, 2, stride=2, scope='layer2-max-pool')
            net = slim.conv2d(net, 40, [5, 5], padding='VALID',
                              scope='layer3-conv')
            net = slim.max_pool2d(net, 2, stride=2, scope='layer4-max-pool')
            net = tf.reshape(net, [-1, 4 * 4 * 40])
            net = slim.fully_connected(net, 256, scope='layer5')
            net = slim.dropout(net, is_training=is_training,
                               scope='layer5-dropout')
            net = slim.fully_connected(net, 256, scope='layer6')
            net = slim.dropout(net, is_training=is_training,
                               scope='layer6-dropout')
            net = slim.fully_connected(net, 10, scope='output',
                                       activation_fn=None)
        return net


def create_experiment(run_config, hparams):
    # You can change a subset of the run_config properties as
    run_config = run_config.replace(save_checkpoints_steps=500)
    # Define the mnist classifier
    estimator = get_estimator(run_config, hparams)
    # Setup data loaders
    train_input_fn, train_input_hook = get_train_inputs(batch_size=128)
    eval_input_fn, eval_input_hook = get_test_inputs(batch_size=128)
    # Define the experiment
    experiment = tf.contrib.learn.Experiment(
        estimator=estimator,
        train_input_fn=train_input_fn,
        eval_input_fn=eval_input_fn,
        train_steps=5000,
        min_eval_frequency=500,
        train_monitors=[train_input_hook],
        eval_hooks=[eval_input_hook],
        eval_steps=None  # Use evaluation feeder until its empty
    )
    return experiment



def main(unused_argv=None):
    model_dir = './mnist_training_new_dataset_hook'

    # Define model parameters
    hparams = tf.contrib.training.HParams(
        learning_rate=0.002,
        n_classes=10
    )

    learn_runner.run(
        experiment_fn=create_experiment,
        run_config=tf.contrib.learn.RunConfig(model_dir=model_dir),
        schedule="train_and_evaluate",
        hparams=hparams)


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

Hi Peter, I ran into the issue of using placeholder in the input function for tf.Estimator and came across your great answer on stackoverflow. It's super helpful! However, there rises another issue with regard to the input function as I try to return a Dataset instead of (next_example, next_label), which is required by the tf.distribute.MirroredStrategy.

Here is function definition:

def get_inputs(input_names, features, labels, shuffle=False, batch_size=64, num_epochs=None, buffer_size=4096, prefetch_buffer_size=2):
    iterator_initializer_hook = IteratorInitializerHook()

    def input_fn():
        target_pl = tf.placeholder(features[0].dtype, shape=features[0].shape, name='targetPlaceholder')
        context_pl = tf.placeholder(features[1].dtype, shape=features[1].shape, name='contextPlaceholder')
        label_pl = tf.placeholder(tf.int32, shape=(None, 1), name='labelPlaceholder')

        dataset = tf.data.Dataset.from_tensor_slices((dict(zip(input_names, [target_pl, context_pl])), label_pl))

        if num_epochs and shuffle:
            dataset = dataset.apply(tf.data.experimental.shuffle_and_repeat(buffer_size, num_epochs))
        elif shuffle:
            dataset = dataset.shuffle(buffer_size)
        elif num_epochs:
            dataset = dataset.repeat(num_epochs)

        dataset = dataset.batch(batch_size=batch_size)

        dataset = dataset.prefetch(buffer_size=prefetch_buffer_size)

        iterator = dataset.make_initializable_iterator()
        x, y = iterator.get_next()

        iterator_initializer_hook.iterator_initializer_func = \
            lambda sess: sess.run(iterator.initializer,
                                  feed_dict={target_pl: features[0],
                                             context_pl: features[1],
                                             label_pl: labels})

        return dataset

    return input_fn, iterator_initializer_hook

The traceback goes like below:

Traceback (most recent call last):
  File "/Users/a1402899/anaconda2/envs/neuralMF_bby/lib/python3.7/site-packages/tensorflow/python/client/session.py", line 1334, in _do_call
    return fn(*args)
  File "/Users/a1402899/anaconda2/envs/neuralMF_bby/lib/python3.7/site-packages/tensorflow/python/client/session.py", line 1319, in _run_fn
    options, feed_dict, fetch_list, target_list, run_metadata)
  File "/Users/a1402899/anaconda2/envs/neuralMF_bby/lib/python3.7/site-packages/tensorflow/python/client/session.py", line 1407, in _call_tf_sessionrun
    run_metadata)
tensorflow.python.framework.errors_impl.InvalidArgumentError: You must feed a value for placeholder tensor 'contextPlaceholder' with dtype int32 and shape [47668]
         [[{{node contextPlaceholder}}]]

Really appreciate it if you would not mind sharing more insights into this. Thanks!

@yzhangswingman I hit this exact problem. Did you manage to solve it?

@wohlbier oh hey, I actually cannot recall at all. Sorry about that. Don’t know if this is helpful but I switched to tf.keras api altogether since TF2 since it seems to be more stable and better maintained.