带参数的Tensorflow serving_input_receiver_fn

Question

I want to add some arguments to the function serving_input_receiver_fn, because the size of the feature array depends of the model. The problem is that the oficial definition of serving_input_receiver_fn is:

serving_input_receiver_fn: A function that takes no argument and returns a ServingInputReceiver. Required for custom models.

My implementation of this function is:

def serving_input_receiver_fn():
    serialized_tf_example = tf.placeholder(dtype=tf.string, shape=[None], name='input_tensors')
    receiver_tensors = {'inputs': serialized_tf_example}
    feature_spec     = {'words': tf.FixedLenFeature([25],tf.int64)}
    features         = tf.parse_example(serialized_tf_example, feature_spec)

    return tf.estimator.export.ServingInputReceiver(features, receiver_tensors)

So, I want that the size ([25]), the name of the features ('words') and the name of the receiver ('inputs') can be variables. There is a chance to have arguments in this function? or another way to do this?

Answer 1

How about using a nested function or closure?

>>> def create_serving_fn(size, feature, inputs):

        def serving_input_receiver_fn():
            serialized_tf_example = tf.placeholder(dtype=tf.string, shape=[None], name='input_tensors')
            receiver_tensors = {inputs: serialized_tf_example}
            feature_spec     = {feature: tf.FixedLenFeature([size],tf.int64)}
            features         = tf.parse_example(serialized_tf_example, feature_spec)    
            return tf.estimator.export.ServingInputReceiver(features, receiver_tensors)
        return serving_input_receiver_fn

    your_serving_fn = create_serving_fn(25, 'words', 'inputs')
    print(your_serving_fn)

<function create_serving_fn.<locals>.serving_input_receiver_fn at 0x7f10df77bf28>

This way, serving_input_receiver_fn has access to arguments passed to create_serving_fn .