Demonstration of keras summary writing bug

keras_summary_repro.py

import glob
import os
import shutil

import numpy
import tensorflow
import pandas as pd
from tensorflow.python.estimator.export.export_output import PredictOutput
from tensorflow.python.saved_model import signature_constants
from protobuf_to_dict import protobuf_to_dict


print("Tensorflow version: {0}".format(tensorflow.VERSION))

DATA_SIZE = 1024
BATCH_SIZE = 32
N_EPOCHS = 1

EMBED_DIM = 100


DATA_DICT = {'input': numpy.random.rand(
    DATA_SIZE, EMBED_DIM).astype('float32')}
LABELS = tensorflow.keras.utils.to_categorical(
    numpy.random.randint(2, size=(DATA_SIZE,)), num_classes=2)

RUNCFG_BASE = {
    'save_summary_steps': 1,
}
KERAS_MODEL_DIR = '/vol/data0/tmp/test_stuff_idunno/tfevent_analyses/keras_summary_repro_keras'
PLAIN_MODEL_DIR = '/vol/data0/tmp/test_stuff_idunno/tfevent_analyses/keras_summary_repro_plain'

for model_dir in [KERAS_MODEL_DIR, PLAIN_MODEL_DIR]:
    if os.path.exists(model_dir):
        shutil.rmtree(model_dir)



def get_input_fn():
    return tensorflow.estimator.inputs.numpy_input_fn(
        x=DATA_DICT, y=LABELS,
        shuffle=True,
        batch_size=BATCH_SIZE,
        num_epochs=N_EPOCHS,
    )


def get_tensors(inp):
    class_layer = tensorflow.keras.layers.Dense(
        2, input_shape=(EMBED_DIM,), activation='softmax', name='classes')(inp)
    tensorflow.summary.scalar('class_norm', tensorflow.norm(class_layer))
    return class_layer


def create_keras_estimator():
    inp = tensorflow.keras.layers.Input(shape=(EMBED_DIM,), name='input')
    class_layer = get_tensors(inp)
    model = tensorflow.keras.models.Model(
        inputs=[inp], outputs=[class_layer])
    model.compile(optimizer='sgd',
                  loss={'classes': 'categorical_crossentropy'},
                  metrics={'classes': 'accuracy'},
                  #target_tensors=tensorflow.get_collection(tensorflow.GraphKeys.SUMMARIES)
                  )

    runcfg = RUNCFG_BASE.copy()
    runcfg['model_dir'] = KERAS_MODEL_DIR
    estimator = tensorflow.keras.estimator.model_to_estimator(
        model, config=tensorflow.estimator.RunConfig(**runcfg))
    return estimator


def create_plain_estimator():
    def model_fn(features, labels, mode, params):
        class_layer = get_tensors(features['input'])
        if mode == tensorflow.estimator.ModeKeys.PREDICT:
            return tensorflow.estimator.EstimatorSpec(
                mode=mode,
                predictions=class_layer,
                export_outputs={
                    signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: PredictOutput(class_layer)
                },
            )

        loss_tensor = tensorflow.losses.softmax_cross_entropy(
            labels, class_layer)
        #tensorflow.summary.scalar('{0}_loss'.format(mode), loss_tensor)

        train_op = tensorflow.contrib.layers.optimize_loss(
            loss=loss_tensor,
            global_step=tensorflow.train.get_global_step(),
            optimizer='SGD',
            learning_rate=0.01,
            summaries=[]
        )

        # since we're already assuming softmax & cross entropy, might as well just hardcode in
        # acc as the eval metric
        # these need to be converted to non-categorical!
        eval_metrics_ops = {
            "accuracy": tensorflow.metrics.accuracy(
                tensorflow.argmax(labels, axis=1),
                tensorflow.argmax(class_layer, axis=1))
        }

        #tensorflow.summary.scalar('accuracy', eval_metrics_ops['accuracy'][1])  # ugly, get the update op

        if mode in (tensorflow.estimator.ModeKeys.TRAIN,
                    tensorflow.estimator.ModeKeys.EVAL):
            return tensorflow.estimator.EstimatorSpec(
                mode=mode,
                loss=loss_tensor,
                train_op=train_op,
                eval_metric_ops=eval_metrics_ops,
            )

    runcfg = RUNCFG_BASE.copy()
    runcfg['model_dir'] = PLAIN_MODEL_DIR

    return tensorflow.estimator.Estimator(
        model_fn,
        config=tensorflow.estimator.RunConfig(**runcfg),
    )


if __name__ == "__main__":
    train_input_fn = get_input_fn()
    eval_input_fn = get_input_fn()

    evalspec = tensorflow.estimator.EvalSpec(eval_input_fn, steps=None,
                                             start_delay_secs=0, throttle_secs=1)
    steps = ((DATA_SIZE // BATCH_SIZE) + 1) * N_EPOCHS
    trainspec = tensorflow.estimator.TrainSpec(train_input_fn, max_steps=steps)

    print("Training with converted Keras model")
    keras_estimator = create_keras_estimator()
    tensorflow.estimator.train_and_evaluate(
        estimator=keras_estimator, train_spec=trainspec, eval_spec=evalspec)

    # get the dang ole file
    keras_events_path = glob.glob(os.path.join(KERAS_MODEL_DIR, 'events.out.tfevents.*'))[0]
    keras_events_df = pd.DataFrame.from_records(map(protobuf_to_dict,
        tensorflow.train.summary_iterator(keras_events_path)))

    print(keras_events_df.summary.dropna().iloc[0])

    print("Training with plain ole Estimator")
    plain_estimator = create_plain_estimator()
    tensorflow.estimator.train_and_evaluate(
        estimator=plain_estimator, train_spec=trainspec, eval_spec=evalspec)

    # get the dang ole file
    plain_events_path = glob.glob(os.path.join(PLAIN_MODEL_DIR, 'events.out.tfevents.*'))[0]
    plain_events_df = pd.DataFrame.from_records(map(protobuf_to_dict,
        tensorflow.train.summary_iterator(plain_events_path)))

    print(plain_events_df.summary.dropna().iloc[0])

In order to run this you'll need some standard stuff (pandas, numpy, tensorflow of course) as well as this weird **protobuf_to_dict utility because I couldn't be arsed to figure out how to filter down the events to summary events without it.

Also obviously change KERAS_MODEL_DIR and PLAIN_MODEL_DIR to something that is valid on your own system.