Tfx: Keras Tutorial produces model that does not integrate with the TF Serving Rest API

[jason-brian-anderson]

Apologies if I am missing something obvious, but i've been struggling with this for a about 5 days and it looks like critical documentation is missing. I'm building a text classifier based off of the Keras TFX tutorial. As far as i can tell, the serving_default signaturedef created by the serve_tf_examples_fn. literally only parses tf.examples.:

@tf.function
  def serve_tf_examples_fn(serialized_tf_examples):
    """Returns the output to be used in the serving signature."""
    feature_spec = tf_transform_output.raw_feature_spec()
    feature_spec.pop(_LABEL_KEY)
    parsed_features = tf.io.parse_example(serialized_tf_examples, feature_spec). #Causes graph to require tf.example input.  what about raw string input?

    transformed_features = model.tft_layer(parsed_features)
    transformed_features.pop(_transformed_name(_LABEL_KEY))

    return model(transformed_features)

This makes perfect sense in terms of TFMA input for the Evaluator component, which works beautifully, but makes no sense in terms of running the model in Serving, as it requires the REST client to import the necessary Tensorflow libraries in order to format and serialize a tf.example for every request as is expected by the line above:

parsed_features = tf.io.parse_example(serialized_tf_examples, feature_spec)

Surely this is not intended as acceptable overhead expected by clients?

I've successfully built an entire TFX pipeline to classify text as a tf.string and am ready to deploy to production but for this one roadblock. I'd be happy to submit a pull request with a documentation update if any experts could point me in the right direction toward documenting the appropriate way to parse in raw text vs a serialized TF.example.

Thank you very much.

[jason-brian-anderson]

Maybe a better way of phrasing this question is: what does a signaturedef look like that would allow a tf transform coupled model to be placed in tf serving so that you can do a simple REST request and not have to create a full serialized tf.example on the client prior to submission?

This source code suggests there is such a thing, but it(transform_raw_features_layer) is nowhere to be found via google search engine!:

  For internal use only.  Users should use the `transform_raw_features` or
  `transform_raw_features_layer` method of the TFTrandformOutput class.

I have a fully functional ML pipeline from 3 weeks with the Keras Tutorial. Once TFX documents this it is going to be completely awesome!

[hanneshapke]

Hi @jason-brian-anderson,

I came across your question. Someone asked a similar question regarding the TFX BERT Pipeline example and therefore, I created an example export tf function to avoid the tf.Example dependency. Maybe the colab example below can be helpful.

Previously, in the Estimator world, we had methods like build_raw_serving_input_receiver_fn to avoid the tf.Example dependency. As far as I understand, we can now define the inputs through the TensorSpecs in the concrete_fuunction() when we export the model. TF is looking for the serve_tf_examples_fn() function but we can remove the tf.Example parsing and add our input handling as shown in the Colab notebook below:

    @tf.function
    def serve_tf_examples_fn(text):
        """Returns the output to be used in the serving signature.
        TF is currently looking for a function serve_tf_examples_fn(),
        therefore the function name can't be changed, even though we aren't 
        parsing tf.Example records.
        """
        reshaped_text = tf.reshape(text, [-1, 1])
        transformed_features = model.tft_layer({"text": reshaped_text})
        outputs = model(transformed_features)
        return {'outputs': outputs}

model.tft_layer expects a dict of the inputs. Because of our input signature of text has the shape [None], we need to reshape the Tensor to match the Keras input requirements.

Colab version of the BERT Pipeline which exports a model for simple REST requests:
https://colab.research.google.com/gist/hanneshapke/f0980b7422d367808dae409536fe9b46/tfx_pipeline_for_bert_preprocessing_wo_tf-example.ipynb

Please note: I am not working for Google and this answer doesn't represent the reply a Tensorflower. It just shows a way of how I solved the same problem.

@jason-brian-anderson Let me know if this solution works for you.

[1025KB]

Hi, @jason-brian-anderson, you can try change the signature to take tensors instead of serialized tf.example see if that works for tf.serving restful api.

here is an example

[ucdmkt]

1108 is likely related

[jason-brian-anderson]

@hanneshapke - first thank you for the help, and second - for the great book that you are co-authoring! I only wish i'd found it before starting this project!

As mentioned, my colleague and I solved the problem in such a way as to 1) allow for the the examplegen component example output (example_gen.outputs['examples']) for model evaluation input to the evaluator component component, while 2) simultaneously providing a signature specifiable in TF Serving context that would allow of raw text predictions. Keeping TFMA validation to prevent the blessing of unworthy models is a key requirement for us.

The following code was added _within the scope_ of the run_fn function. (note that instead of 'text' as an input, our input is 'request_url':

    class MyModule(tf.Module):
        def __init__(self, model, tf_transform_output):
            self.model = model
            self.tf_transform_output = tf_transform_output
            self.model.tft_layer = self.tf_transform_output.transform_features_layer()

        @tf.function(input_signature=[tf.TensorSpec(shape=(None), dtype=tf.string, name='request_url')])
        def tf_serving_raw_input_fn(self, request_url):

            request_url_sp_tensor = tf.sparse.SparseTensor(
                indices=[[0, 0]],
                values=request_url,
                dense_shape=(1, 1)
            )

            parsed_features = {'request_url': request_url_sp_tensor, }
            transformed_features = self.model.tft_layer(parsed_features)
            transformed_features.pop(_transformed_name(_LABEL_KEY))
            return self.model(transformed_features)

        @tf.function(input_signature=[tf.TensorSpec(shape=[None], dtype=tf.string, name='examples')])
        def serve_tf_examples_fn(self, serialized_tf_examples):
            feature_spec = self.tf_transform_output.raw_feature_spec()
            feature_spec.pop(_LABEL_KEY)
            parsed_features = tf.io.parse_example(serialized_tf_examples, feature_spec)
            transformed_features = self.model.tft_layer(parsed_features)
            transformed_features.pop(_transformed_name(_LABEL_KEY))
            return self.model(transformed_features)


    module = MyModule(model, tf_transform_output)

    signatures = {"serving_default": module.serve_tf_examples_fn,
                  "serving_raw_input": module.tf_serving_raw_input_fn,

                  }
    tf.saved_model.save(module,
                        export_dir=fn_args.serving_model_dir,
                        signatures=signatures,
                        options=None,
                        )
                        )

The first obvious change is subclassing tf.module instead of directly saving the Keras model as is done in the TF2 Keras Tutorial. This approach was inspired by this stackoverflow,but it doesn't explain the difference in terms of this approach vs directly saving the keras model with the get_concrete_function approach.I'm not sure whether using this approach contributed to our success or whether a direct keras save with the above tf.function would have also worked.

The REST API code was as follows:
````
rest_uri = f'http://{server}:8501/v1/models/wrcv3/versions/{ver}:predict'
headers = {"content-type": "application/json"}

data = {
"signature_name":"serving_raw_input",
"instances":[
{
"request_url":"cgi-bin"
},

]
}
data = json.dumps(data)
json_response = requests.post(rest_uri, data=data, headers=headers)
print(json_response.content)
print(json_response.json)
resulting in:
b'{\n "predictions": [[0.997684]\n ]\n}'
````

One notable drawback to our solution above it that it is not capable of receiving multiple inputs. Although we expect this to be an easy fix. Not sure if this will affect TF Serving's batching ability. Haven't tested that yet.

I did carefully read both your response and spent time with the referenced colab, Referring to your cited code comment:

@tf.function def serve_tf_examples_fn(text): """Returns the output to be used in the serving signature. TF is currently looking for a function serve_tf_examples_fn(), therefore the function name can't be changed, even though we aren't parsing tf.Example records. """ reshaped_text = tf.reshape(text, [-1, 1]) transformed_features = model.tft_layer({"text": reshaped_text}) outputs = model(transformed_features) return {'outputs': outputs}
This presents a problem for us in that Evaluator component by default uses the serving_default signature which our tf2 keras tutorial -based serve_tf_examples_fn is keyed to in the signatures dict.

To test the assertion made above, i renamed serve_tf_examples_fn to alternatively_named_serve_tf_examples_fn as an experiment. This did result in a fully successful Airflow DAG along with a BLESSED model, TFMA validation metrics, and several render_slicing_metrics output from TFMA. So i'm wondering what the discrepency is between what you have observed and this?

Regarding the specific content of the serve_tf_examples_fn you cite above, using the tf.reshape seems like a great way to handle an incoming dense tensor, but as we stuck very closely with the tutorial's code, we use the _fill_in_missing as the first operation in our preprocessing_fn function, which expects and parses sparsetensors, and not dense tensors. We did try it anyway but it didn't work, the traceback suggesting that tf.reshape was not happy about receiving a sparse tensor.
Perhaps our use of the tf.Module subclassing approach enabled us to evade the constraint you found above? and if so, then perhaps that constraint is limited to the keras.save functionality, or the tf.saved_model.save() code on which (i think) the keras save functionality depends?

We are new to TF2 and TFX so if we've misunderstood or mis-stated anything please feel encouraged to correct us!

Side note, the signaturedefs documentation on the tfx site that lists the predict/classify/regression protobufs is still an absolute mystery to us - in terms of how to use them at least.

[tensorflow-butler[bot]]

Are you satisfied with the resolution of your issue?
Yes
No

[1025KB]

HI, Did you try change the signature to take tensor like this

In this way the signature can take tensors and it works with tfma.

[C45513]

Hi @jason-brian-anderson, thanks for your details comment above. I'm having the same issue and I got the following error with the concept that you had. I wonder if you can help me with that. Thanks in advance! Instead of request URL, I have multiple data features as tf.string.

Error:

    110           ("Expected a TensorFlow function to generate a signature for, but "
    111            "got {}. Only `tf.functions` with an input signature or "
--> 112            "concrete functions can be used as a signature.").format(function))
    113 
    114     wrapped_functions[original_function] = signature_function = (

ValueError: Expected a TensorFlow function to generate a signature for, but got . Only `tf.functions` with an input signature or concrete functions can be used as a signature.

Code:

def run_fn(fn_args):

    tf_transform_output = tft.TFTransformOutput(fn_args.transform_output)

    train_dataset = _input_fn(fn_args.train_files, tf_transform_output, 40)
    eval_dataset = _input_fn(fn_args.eval_files, tf_transform_output, 40)

    model = get_model()

    model.fit(
      train_dataset,
      steps_per_epoch=fn_args.train_steps,
      validation_data=eval_dataset,
      validation_steps=fn_args.eval_steps)

    class MyModule(tf.Module):
        def __init__(self, model, tf_transform_output):
            self.model = get_model
            self.tf_transform_output = tf_transform_output
            self.model.tft_layer = self.tf_transform_output.transform_features_layer()
        @tf.function
        def serve_tf_examples_fn(serialized_tf_examples):
            """Returns the output to be used in the serving signature."""
            feature_spec = tf_transform_output.raw_feature_spec()
            feature_spec.pop(_LABEL_KEY_EA)
            parsed_features = tf.io.parse_example(serialized_tf_examples, feature_spec)

            transformed_features = model.tft_layer(parsed_features)
            transformed_features.pop(_transformed_name(_LABEL_KEY_EA))
            return self.model(transformed_features)

        @tf.function(input_signature=[tf.TensorSpec(shape=(None), dtype=tf.string, name='raw_examples')])
        def tf_serving_raw_input_fn(self, raw_features):
            raw_features_sp_tensor = tf.sparse.SparseTensor(
                indices=[[0, 0]],
                values=raw_features,
                dense_shape=(1, 1)
            )
            parsed_features = {
                'raw_features': raw_features_sp_tensor,
            }
            transformed_features = self.model.tft_example_layer(parsed_features)
            transformed_features.pop(_transformed_name(_LABEL_KEY_EA))
            return self.model(transformed_features)

    module = MyModule(model, tf_transform_output)

    signatures = {
      'serving_default': module.serve_tf_examples_fn,
      'serving_raw_input': module.tf_serving_raw_input_fn
    }

    tf.saved_model.save(module, export_dir=fn_args.serving_model_dir, signatures=signatures, options=None,)