Keras estimator bug repro

endtoend_estimator_keras_repro.py

from __future__ import division

import os

import numpy
import tensorflow
from tensorflow.python.keras._impl import keras
from tensorflow.python.estimator.export.export_output import PredictOutput
from tensorflow.python.estimator.export.export import build_raw_serving_input_receiver_fn
from tensorflow.python.saved_model import signature_constants
from tensorflow.python import debug as tf_debug

#import datautil
#import ml_utils
# still working out how to best do this:

# ENDTOEND embedding only text model

# I have no clue what I'm doing
# reference: https://github.com/awslabs/amazon-sagemaker-examples/blob/master/sagemaker-python-sdk/tensorflow_abalone_age_predictor_using_keras/abalone.py

EMB_DIM = 10
EMB_SIZE = 4

NUM_SEQ = 6
SEQ_LEN = 3

INP_NAME = 'lookedup'


# fuck it
def simple_keras_model(inp_placeholder, embedding_mat):
    # just an embedding layer
    #emb = keras.layers.Embedding(*(embedding_mat.shape), weights=[embedding_mat],
    #                             input_length=SEQ_LEN, name='embed', trainable=False)(inp_placeholder)
    emb = keras.layers.Embedding(*(embedding_mat.shape), embeddings_initializer=keras.initializers.Constant(value=embedding_mat),
                                 input_length=SEQ_LEN, name='embed', trainable=False)(inp_placeholder)

    pooler = keras.layers.GlobalAveragePooling1D(name='avg')(emb)
    model = keras.layers.core.Dense(1, input_shape=(EMB_DIM,), name='weights', use_bias=False,)(pooler)
    return {'emb':emb, 'pooler':pooler, 'model':model}

class model_fn_defaults:
    train_spec = {"optimizer":"SGD", "lr":0.00}
    compile_spec = {
            "optimizer":"sgd",
            "loss":{"weights":"mean_squared_error"},
    }
    inp = INP_NAME

def model_fn(features, labels, mode, params):
    train_spec = params.get('train_spec', model_fn_defaults.train_spec)

    inp = params.get('inp', model_fn_defaults.inp)

    # ensure that all keras stuff is done on the default graph
    #keras.backend.set_session(
    #        tensorflow.Session(
    #            graph=tensorflow.get_default_graph()
    #        )
    #)
    with tensorflow.Session(graph=tensorflow.get_default_graph()) as sess:
        #keras.backend.set_session(sess)
        tensorflow.set_random_seed(4)
        def_graph = tensorflow.get_default_graph()
        print(def_graph)
        keras_graph = keras.backend.get_session().graph
        print(keras_graph)
        inp_placeholder = features[inp]

        basic_emb_mat = numpy.random.randint(EMB_SIZE, size=(EMB_SIZE, EMB_DIM))
        tensors = simple_keras_model(inp_placeholder, basic_emb_mat)

        pred_dict = tensors

        if mode == tensorflow.estimator.ModeKeys.PREDICT:
            return tensorflow.estimator.EstimatorSpec(
                        mode=mode,
                        predictions=pred_dict,
                        export_outputs={
                            signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:PredictOutput(pred_dict)
                        },
                    )

        loss_tensor = tensorflow.nn.l2_loss(tensors['model'], name='l2_loss')
        tensorflow.summary.scalar('l2_loss', loss_tensor)

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

        eval_metrics_ops = {
                "zero":tensorflow.metrics.mean(tensors['model'])
        }

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

        if mode in (tensorflow.estimator.ModeKeys.TRAIN,
                    tensorflow.estimator.ModeKeys.EVAL):
            # put in training and eval ops
            return tensorflow.estimator.EstimatorSpec(
                       mode=mode,
                       loss=loss_tensor,
                       train_op=train_op,
                       eval_metric_ops=eval_metrics_ops
                   )

def keras_model_fn(params):
    compile_spec = params.get('compile_spec', model_fn_defaults.compile_spec)

    inp = keras.layers.Input(shape=(SEQ_LEN,), name=INP_NAME, dtype='int32')
    basic_emb_mat = numpy.random.randint(EMB_SIZE, size=(EMB_SIZE, EMB_DIM))
    tensors = simple_keras_model(inp, basic_emb_mat)
    model = keras.models.Model(inputs=[inp], outputs=[tensors['model']])
    model.compile(**compile_spec)
    return model



def _main_input_fn(inp_name, num_epochs, num_seq=NUM_SEQ):
    numpy.random.seed(4)

    data = numpy.random.randint(EMB_SIZE, size=(num_seq, SEQ_LEN))

    data_dict = {inp_name:data}
    labels = numpy.zeros(num_seq).astype(numpy.float32)
    return tensorflow.estimator.inputs.numpy_input_fn(
            x=data_dict,
            y=labels,
            shuffle=False,
            batch_size=32,
            num_epochs=num_epochs
    )


class train_input_fn_defaults:
    inp_name = INP_NAME
    num_epochs = None # train by number of steps by default
    num_seq = 1

def train_input_fn(params):
    inp_name = params.get('inp_name', train_input_fn_defaults.inp_name)
    num_epochs = params.get('train_num_epochs', train_input_fn_defaults.num_epochs)
    num_seq = params.get('train_num_seq', train_input_fn_defaults.num_seq)

    return _main_input_fn(inp_name, num_epochs, num_seq=num_seq)

class eval_input_fn_defaults:
    inp_name = INP_NAME
    num_epochs = 1
    num_seq = NUM_SEQ

def eval_input_fn(params):
    inp_name = params.get('inp_name', eval_input_fn_defaults.inp_name)
    num_epochs = params.get('eval_num_epochs', eval_input_fn_defaults.num_epochs)
    num_seq = params.get('eval_num_seq', eval_input_fn_defaults.num_seq)

    return _main_input_fn(inp_name, num_epochs, num_seq=num_seq)

class serving_input_fn_defaults:
    tensor_name = INP_NAME

def serving_input_fn(params):
    tensor_name = params.get('tensor_name', serving_input_fn_defaults.tensor_name)
    # not sure what the hyperparameters are for
    # annoyingly I can't really give this the model and say, hey, get what you need
    # as I did in the previous stuff?
    tensor = tensorflow.placeholder(tensorflow.int32, shape=[None,SEQ_LEN])
    return build_raw_serving_input_receiver_fn({
        tensor_name:tensor
    })


if __name__ == "__main__":
    import sys
    print("Testing {0} locally".format(sys.argv[0]))

    hyperparameters = {
                "eval_num_epochs":1,
                "train_num_epochs":0,
    }

    print("Making estimator")


    #estimator = tensorflow.estimator.Estimator(model_fn=model_fn, params=hyperparameters)
    model = keras_model_fn(hyperparameters)
    estimator = keras.estimator.model_to_estimator(keras_model=model)
    print("Training estimator")

    estimator.train(input_fn=train_input_fn(hyperparameters), steps=None,)
            #hooks=[tf_debug.LocalCLIDebugHook()])
    print("Evaluating estimator")
    eval_dict = estimator.evaluate(input_fn=eval_input_fn(hyperparameters), steps=None,)
            #hooks=[tf_debug.LocalCLIDebugHook()])
    print(eval_dict)

    print("Exporting keras_estimator")
    estimator.export_savedmodel(
            os.path.join(estimator.model_dir, 'exported'),
            serving_input_fn(hyperparameters),
    )

    from tensorflow.python.tools import inspect_checkpoint
    print("serialized model")
    inspect_checkpoint.print_tensors_in_checkpoint_file(
            os.path.join(
                estimator.model_dir,
                'model.ckpt-0'),
            "",
            True
    )


    print("Predicting")
    predicted_1 = estimator.predict(input_fn=eval_input_fn(hyperparameters),
        hooks=[tf_debug.LocalCLIDebugHook()]
    )
    model_outputs_1 = []
    for output in predicted_1:
        print(output)
        model_outputs_1.append(output['model'])
    print(numpy.sum(model_outputs_1))