keras export outputs issue

endtoend_estimator_keras_repro_2.py

from __future__ import division

import os

import numpy
import tensorflow as tf
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

EMB_DIM = 10
EMB_SIZE = 4

NUM_SEQ = 6
SEQ_LEN = 3

INP_NAME = 'lookedup'


def simple_keras_model(inp_placeholder, embedding_mat):
    emb = tf.keras.layers.Embedding(*(embedding_mat.shape), embeddings_initializer=tf.keras.initializers.Constant(value=embedding_mat),
                                 input_length=SEQ_LEN, name='embed', trainable=False)(inp_placeholder)

    pooler = tf.keras.layers.GlobalAveragePooling1D(name='avg')(emb)
    model = tf.keras.layers.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 keras_model_fn(params):
    compile_spec = params.get('compile_spec', model_fn_defaults.compile_spec)

    inp = tf.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 = tf.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 tf.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)
    tensor = tf.placeholder(tf.int32, shape=[None,SEQ_LEN])
    return tf.estimator.export.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")


    model = keras_model_fn(hyperparameters)
    estimator = tf.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 tf.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))