TFX with AI Platform Serving

[luischinchillagarcia]

When using AI Platform for serving a model, the model requires a JSON in the form {"instances": [...] }.

However, when creating a TFX pipeline, the model only uses Examples as inputs and outputs. This means that in order to get predictions we have to turn every raw input into Examples, turn the Examples into a JSON, and pass them through the model via the AI Platform API client (and also extract the predictions from the output model Examples).

Is there a way to integrate TFX with AI Platform serving where the transformation from raw data to Examples is done internally or the optimal way to do this integration?

[ucdmkt]

This is because the default demonstration of Trainer code (the equivalent of taxi_utils.py), probably yours as well, is implemented in such a way that it exports the model that expects 1-D bytes Tensor with serialized tf.Examples protocol buffer [code].

If you export the model with a serving_input_receiver_fn() in a way that receives flat list of named Tensors as input [for example, code], AI Platform prediction can receive input Tensors as flat list of JSON arrays.

[rmothukuru]

@luischinchillagarcia,
Can you please respond to @ucdmkt's comment. Thanks!

[luischinchillagarcia]

Yes! Thanks for the quick response, and apologies, I was working through another (similar) issue I kept encountering.

This definitely solved the issue and will close this issue. Thanks!

[C45513]

Hi, I'm new to tensforflow as well as tfx and I ran into the same issue when the model is being served on Google AI Platform Model as part of the tfx pipeline, but the JSON doesn't accept the format. I'm following the keras tutorial the data input has to be serialized as follow code.

@ucdmkt ,I've tried with the code sample that you posted, but I can't seem to get it to work. Can you please help? With the schema name error I got, do you mean to refer to the outputs of schema_gen? FYI, the data types of the data features are tf.string

Error:

    111     signatures = {
    112       'serving_default':
--> 113           _get_serve_tf_examples_fn(tf_transform_output, schema).get_concrete_function(
    114                                         tf.TensorSpec(
    115                                             shape=[None],

NameError: name 'schema' is not defined

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)

    signatures = {
      'serving_default':
          _get_serve_tf_examples_fn(tf_transform_output, schema).get_concrete_function(
                                        tf.TensorSpec(
                                            shape=[None],
                                            dtype=tf.string,
                                            name='examples')),
    }
    model.save(fn_args.serving_model_dir, save_format='tf', signatures=signatures)
    
def _input_fn(file_pattern,
              tf_transform_output,
              batch_size= 32):

    transformed_feature_spec = (
      tf_transform_output.transformed_feature_spec().copy())

    dataset = tf.data.experimental.make_batched_features_dataset(
      file_pattern=file_pattern,
      batch_size=batch_size,
      features=transformed_feature_spec,
      reader=_gzip_reader_fn,
      label_key=_transformed_name(_LABEL_KEY_EA))

    return dataset

def _get_raw_feature_spec(schema):
  return schema_utils.schema_as_feature_spec(schema).feature_spec

def _get_serve_tf_examples_fn(tf_transform_output, schema):
    """Build the serving function for flat list of Dense tensors as input.
    Args:
        tf_transform_output: A TFTransformOutput.
        schema: the schema of the input data.
    Returns:
        Tensorflow graph which parses examples, applying tf-transform to them.
    """
    raw_feature_spec = _get_raw_feature_spec(schema)
    raw_feature_spec.pop(_LABEL_KEY_EA)

    raw_input_fn = tf.estimator.export.build_parsing_serving_input_receiver_fn(
        raw_feature_spec, default_batch_size=None)
    serving_input_receiver = raw_input_fn()

    transformed_features = tf_transform_output.transform_raw_features(
        serving_input_receiver.features)

    return tf.estimator.export.ServingInputReceiver(
        transformed_features, serving_input_receiver.features)

[snehpandya123]

@C45513 Were you able to solve above problem?