Copyright 2021 The TensorFlow Authors.¶

In [ ]:

#@title Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

#@title Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # https://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License.

Feature Engineering using TFX Pipeline and TensorFlow Transform¶

Transform input data and train a model with a TFX pipeline.

Note: We recommend running this tutorial in a Colab notebook, with no setup required! Just click "Run in Google Colab".

View on TensorFlow.org

Run in Google Colab

View source on GitHub

Download notebook

In this notebook-based tutorial, we will create and run a TFX pipeline to ingest raw input data and preprocess it appropriately for ML training. This notebook is based on the TFX pipeline we built in Data validation using TFX Pipeline and TensorFlow Data Validation Tutorial. If you have not read that one yet, you should read it before proceeding with this notebook.

You can increase the predictive quality of your data and/or reduce dimensionality with feature engineering. One of the benefits of using TFX is that you will write your transformation code once, and the resulting transforms will be consistent between training and serving in order to avoid training/serving skew.

We will add a Transform component to the pipeline. The Transform component is implemented using the tf.transform library.

Please see Understanding TFX Pipelines to learn more about various concepts in TFX.

Set Up¶

We first need to install the TFX Python package and download the dataset which we will use for our model.

Upgrade Pip¶

To avoid upgrading Pip in a system when running locally, check to make sure that we are running in Colab. Local systems can of course be upgraded separately.

In [ ]:

try:
  import colab
  !pip install --upgrade pip
except:
  pass

try: import colab !pip install --upgrade pip except: pass

Install TFX¶

In [ ]:

!pip install -U tfx

!pip install -U tfx

Did you restart the runtime?¶

If you are using Google Colab, the first time that you run the cell above, you must restart the runtime by clicking above "RESTART RUNTIME" button or using "Runtime > Restart runtime ..." menu. This is because of the way that Colab loads packages.

Check the TensorFlow and TFX versions.

In [ ]:

import tensorflow as tf
print('TensorFlow version: {}'.format(tf.__version__))
from tfx import v1 as tfx
print('TFX version: {}'.format(tfx.__version__))

import tensorflow as tf print('TensorFlow version: {}'.format(tf.__version__)) from tfx import v1 as tfx print('TFX version: {}'.format(tfx.__version__))

Set up variables¶

There are some variables used to define a pipeline. You can customize these variables as you want. By default all output from the pipeline will be generated under the current directory.

In [ ]:

import os

PIPELINE_NAME = "penguin-transform"

# Output directory to store artifacts generated from the pipeline.
PIPELINE_ROOT = os.path.join('pipelines', PIPELINE_NAME)
# Path to a SQLite DB file to use as an MLMD storage.
METADATA_PATH = os.path.join('metadata', PIPELINE_NAME, 'metadata.db')
# Output directory where created models from the pipeline will be exported.
SERVING_MODEL_DIR = os.path.join('serving_model', PIPELINE_NAME)

from absl import logging
logging.set_verbosity(logging.INFO)  # Set default logging level.

import os PIPELINE_NAME = "penguin-transform" # Output directory to store artifacts generated from the pipeline. PIPELINE_ROOT = os.path.join('pipelines', PIPELINE_NAME) # Path to a SQLite DB file to use as an MLMD storage. METADATA_PATH = os.path.join('metadata', PIPELINE_NAME, 'metadata.db') # Output directory where created models from the pipeline will be exported. SERVING_MODEL_DIR = os.path.join('serving_model', PIPELINE_NAME) from absl import logging logging.set_verbosity(logging.INFO) # Set default logging level.

Prepare example data¶

We will download the example dataset for use in our TFX pipeline. The dataset we are using is Palmer Penguins dataset.

However, unlike previous tutorials which used an already preprocessed dataset, we will use the raw Palmer Penguins dataset.

Because the TFX ExampleGen component reads inputs from a directory, we need to create a directory and copy the dataset to it.

In [ ]:

import urllib.request
import tempfile

DATA_ROOT = tempfile.mkdtemp(prefix='tfx-data')  # Create a temporary directory.
_data_path = 'https://storage.googleapis.com/download.tensorflow.org/data/palmer_penguins/penguins_size.csv'
_data_filepath = os.path.join(DATA_ROOT, "data.csv")
urllib.request.urlretrieve(_data_path, _data_filepath)

import urllib.request import tempfile DATA_ROOT = tempfile.mkdtemp(prefix='tfx-data') # Create a temporary directory. _data_path = 'https://storage.googleapis.com/download.tensorflow.org/data/palmer_penguins/penguins_size.csv' _data_filepath = os.path.join(DATA_ROOT, "data.csv") urllib.request.urlretrieve(_data_path, _data_filepath)

Take a quick look at what the raw data looks like.

In [ ]:

!head {_data_filepath}

!head {_data_filepath}

There are some entries with missing values which are represented as NA. We will just delete those entries in this tutorial.

In [ ]:

!sed -i '/\bNA\b/d' {_data_filepath}
!head {_data_filepath}

!sed -i '/\bNA\b/d' {_data_filepath} !head {_data_filepath}

You should be able to see seven features which describe penguins. We will use the same set of features as the previous tutorials - 'culmen_length_mm', 'culmen_depth_mm', 'flipper_length_mm', 'body_mass_g' - and will predict the 'species' of a penguin.

The only difference will be that the input data is not preprocessed. Note that we will not use other features like 'island' or 'sex' in this tutorial.

Prepare a schema file¶

As described in Data validation using TFX Pipeline and TensorFlow Data Validation Tutorial, we need a schema file for the dataset. Because the dataset is different from the previous tutorial we need to generate it again. In this tutorial, we will skip those steps and just use a prepared schema file.

In [ ]:

import shutil

SCHEMA_PATH = 'schema'

_schema_uri = 'https://raw.githubusercontent.com/tensorflow/tfx/master/tfx/examples/penguin/schema/raw/schema.pbtxt'
_schema_filename = 'schema.pbtxt'
_schema_filepath = os.path.join(SCHEMA_PATH, _schema_filename)

os.makedirs(SCHEMA_PATH, exist_ok=True)
urllib.request.urlretrieve(_schema_uri, _schema_filepath)

import shutil SCHEMA_PATH = 'schema' _schema_uri = 'https://raw.githubusercontent.com/tensorflow/tfx/master/tfx/examples/penguin/schema/raw/schema.pbtxt' _schema_filename = 'schema.pbtxt' _schema_filepath = os.path.join(SCHEMA_PATH, _schema_filename) os.makedirs(SCHEMA_PATH, exist_ok=True) urllib.request.urlretrieve(_schema_uri, _schema_filepath)

This schema file was created with the same pipeline as in the previous tutorial without any manual changes.

Create a pipeline¶

TFX pipelines are defined using Python APIs. We will add Transform component to the pipeline we created in the Data Validation tutorial.

A Transform component requires input data from an ExampleGen component and a schema from a SchemaGen component, and produces a "transform graph". The output will be used in a Trainer component. Transform can optionally produce "transformed data" in addition, which is the materialized data after transformation. However, we will transform data during training in this tutorial without materialization of the intermediate transformed data.

One thing to note is that we need to define a Python function, preprocessing_fn to describe how input data should be transformed. This is similar to a Trainer component which also requires user code for model definition.

Write preprocessing and training code¶

We need to define two Python functions. One for Transform and one for Trainer.

preprocessing_fn¶

The Transform component will find a function named preprocessing_fn in the given module file as we did for Trainer component. You can also specify a specific function using the preprocessing_fn parameter of the Transform component.

In this example, we will do two kinds of transformation. For continuous numeric features like culmen_length_mm and body_mass_g, we will normalize these values using the tft.scale_to_z_score function. For the label feature, we need to convert string labels into numeric index values. We will use tf.lookup.StaticHashTable for conversion.

run_fn¶

The model itself is almost the same as in the previous tutorials, but this time we will transform the input data using the transform graph from the Transform component.

One more important difference compared to the previous tutorial is that now we export a model for serving which includes not only the computation graph of the model, but also the transform graph for preprocessing, which is generated in Transform component. We need to define a separate function which will be used for serving incoming requests. You can see that the same function _apply_preprocessing was used for both of the training data and the serving request.

In [ ]:

_module_file = 'penguin_utils.py'

_module_file = 'penguin_utils.py'

In [ ]:

%%writefile {_module_file}


from typing import List, Text
from absl import logging
import tensorflow as tf
from tensorflow import keras
from tensorflow_metadata.proto.v0 import schema_pb2
import tensorflow_transform as tft
from tensorflow_transform.tf_metadata import schema_utils

from tfx import v1 as tfx
from tfx_bsl.public import tfxio

# Specify features that we will use.
_FEATURE_KEYS = [
    'culmen_length_mm', 'culmen_depth_mm', 'flipper_length_mm', 'body_mass_g'
]
_LABEL_KEY = 'species'

_TRAIN_BATCH_SIZE = 20
_EVAL_BATCH_SIZE = 10


# NEW: TFX Transform will call this function.
def preprocessing_fn(inputs):
  """tf.transform's callback function for preprocessing inputs.

  Args:
    inputs: map from feature keys to raw not-yet-transformed features.

  Returns:
    Map from string feature key to transformed feature.
  """
  outputs = {}

  # Uses features defined in _FEATURE_KEYS only.
  for key in _FEATURE_KEYS:
    # tft.scale_to_z_score computes the mean and variance of the given feature
    # and scales the output based on the result.
    outputs[key] = tft.scale_to_z_score(inputs[key])

  # For the label column we provide the mapping from string to index.
  # We could instead use `tft.compute_and_apply_vocabulary()` in order to
  # compute the vocabulary dynamically and perform a lookup.
  # Since in this example there are only 3 possible values, we use a hard-coded
  # table for simplicity.
  table_keys = ['Adelie', 'Chinstrap', 'Gentoo']
  initializer = tf.lookup.KeyValueTensorInitializer(
      keys=table_keys,
      values=tf.cast(tf.range(len(table_keys)), tf.int64),
      key_dtype=tf.string,
      value_dtype=tf.int64)
  table = tf.lookup.StaticHashTable(initializer, default_value=-1)
  outputs[_LABEL_KEY] = table.lookup(inputs[_LABEL_KEY])

  return outputs


# NEW: This function will apply the same transform operation to training data
#      and serving requests.
def _apply_preprocessing(raw_features, tft_layer):
  transformed_features = tft_layer(raw_features)
  if _LABEL_KEY in raw_features:
    transformed_label = transformed_features.pop(_LABEL_KEY)
    return transformed_features, transformed_label
  else:
    return transformed_features, None


# NEW: This function will create a handler function which gets a serialized
#      tf.example, preprocess and run an inference with it.
def _get_serve_tf_examples_fn(model, tf_transform_output):
  # We must save the tft_layer to the model to ensure its assets are kept and
  # tracked.
  model.tft_layer = tf_transform_output.transform_features_layer()

  @tf.function(input_signature=[
      tf.TensorSpec(shape=[None], dtype=tf.string, name='examples')
  ])
  def serve_tf_examples_fn(serialized_tf_examples):
    # Expected input is a string which is serialized tf.Example format.
    feature_spec = tf_transform_output.raw_feature_spec()
    # Because input schema includes unnecessary fields like 'species' and
    # 'island', we filter feature_spec to include required keys only.
    required_feature_spec = {
        k: v for k, v in feature_spec.items() if k in _FEATURE_KEYS
    }
    parsed_features = tf.io.parse_example(serialized_tf_examples,
                                          required_feature_spec)

    # Preprocess parsed input with transform operation defined in
    # preprocessing_fn().
    transformed_features, _ = _apply_preprocessing(parsed_features,
                                                   model.tft_layer)
    # Run inference with ML model.
    return model(transformed_features)

  return serve_tf_examples_fn


def _input_fn(file_pattern: List[Text],
              data_accessor: tfx.components.DataAccessor,
              tf_transform_output: tft.TFTransformOutput,
              batch_size: int = 200) -> tf.data.Dataset:
  """Generates features and label for tuning/training.

  Args:
    file_pattern: List of paths or patterns of input tfrecord files.
    data_accessor: DataAccessor for converting input to RecordBatch.
    tf_transform_output: A TFTransformOutput.
    batch_size: representing the number of consecutive elements of returned
      dataset to combine in a single batch

  Returns:
    A dataset that contains (features, indices) tuple where features is a
      dictionary of Tensors, and indices is a single Tensor of label indices.
  """
  dataset = data_accessor.tf_dataset_factory(
      file_pattern,
      tfxio.TensorFlowDatasetOptions(batch_size=batch_size),
      schema=tf_transform_output.raw_metadata.schema)

  transform_layer = tf_transform_output.transform_features_layer()
  def apply_transform(raw_features):
    return _apply_preprocessing(raw_features, transform_layer)

  return dataset.map(apply_transform).repeat()


def _build_keras_model() -> tf.keras.Model:
  """Creates a DNN Keras model for classifying penguin data.

  Returns:
    A Keras Model.
  """
  # The model below is built with Functional API, please refer to
  # https://www.tensorflow.org/guide/keras/overview for all API options.
  inputs = [
      keras.layers.Input(shape=(1,), name=key)
      for key in _FEATURE_KEYS
  ]
  d = keras.layers.concatenate(inputs)
  for _ in range(2):
    d = keras.layers.Dense(8, activation='relu')(d)
  outputs = keras.layers.Dense(3)(d)

  model = keras.Model(inputs=inputs, outputs=outputs)
  model.compile(
      optimizer=keras.optimizers.Adam(1e-2),
      loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
      metrics=[keras.metrics.SparseCategoricalAccuracy()])

  model.summary(print_fn=logging.info)
  return model


# TFX Trainer will call this function.
def run_fn(fn_args: tfx.components.FnArgs):
  """Train the model based on given args.

  Args:
    fn_args: Holds args used to train the model as name/value pairs.
  """
  tf_transform_output = tft.TFTransformOutput(fn_args.transform_output)

  train_dataset = _input_fn(
      fn_args.train_files,
      fn_args.data_accessor,
      tf_transform_output,
      batch_size=_TRAIN_BATCH_SIZE)
  eval_dataset = _input_fn(
      fn_args.eval_files,
      fn_args.data_accessor,
      tf_transform_output,
      batch_size=_EVAL_BATCH_SIZE)

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

  # NEW: Save a computation graph including transform layer.
  signatures = {
      'serving_default': _get_serve_tf_examples_fn(model, tf_transform_output),
  }
  model.save(fn_args.serving_model_dir, save_format='tf', signatures=signatures)

%%writefile {_module_file} from typing import List, Text from absl import logging import tensorflow as tf from tensorflow import keras from tensorflow_metadata.proto.v0 import schema_pb2 import tensorflow_transform as tft from tensorflow_transform.tf_metadata import schema_utils from tfx import v1 as tfx from tfx_bsl.public import tfxio # Specify features that we will use. _FEATURE_KEYS = [ 'culmen_length_mm', 'culmen_depth_mm', 'flipper_length_mm', 'body_mass_g' ] _LABEL_KEY = 'species' _TRAIN_BATCH_SIZE = 20 _EVAL_BATCH_SIZE = 10 # NEW: TFX Transform will call this function. def preprocessing_fn(inputs): """tf.transform's callback function for preprocessing inputs. Args: inputs: map from feature keys to raw not-yet-transformed features. Returns: Map from string feature key to transformed feature. """ outputs = {} # Uses features defined in _FEATURE_KEYS only. for key in _FEATURE_KEYS: # tft.scale_to_z_score computes the mean and variance of the given feature # and scales the output based on the result. outputs[key] = tft.scale_to_z_score(inputs[key]) # For the label column we provide the mapping from string to index. # We could instead use `tft.compute_and_apply_vocabulary()` in order to # compute the vocabulary dynamically and perform a lookup. # Since in this example there are only 3 possible values, we use a hard-coded # table for simplicity. table_keys = ['Adelie', 'Chinstrap', 'Gentoo'] initializer = tf.lookup.KeyValueTensorInitializer( keys=table_keys, values=tf.cast(tf.range(len(table_keys)), tf.int64), key_dtype=tf.string, value_dtype=tf.int64) table = tf.lookup.StaticHashTable(initializer, default_value=-1) outputs[_LABEL_KEY] = table.lookup(inputs[_LABEL_KEY]) return outputs # NEW: This function will apply the same transform operation to training data # and serving requests. def _apply_preprocessing(raw_features, tft_layer): transformed_features = tft_layer(raw_features) if _LABEL_KEY in raw_features: transformed_label = transformed_features.pop(_LABEL_KEY) return transformed_features, transformed_label else: return transformed_features, None # NEW: This function will create a handler function which gets a serialized # tf.example, preprocess and run an inference with it. def _get_serve_tf_examples_fn(model, tf_transform_output): # We must save the tft_layer to the model to ensure its assets are kept and # tracked. model.tft_layer = tf_transform_output.transform_features_layer() @tf.function(input_signature=[ tf.TensorSpec(shape=[None], dtype=tf.string, name='examples') ]) def serve_tf_examples_fn(serialized_tf_examples): # Expected input is a string which is serialized tf.Example format. feature_spec = tf_transform_output.raw_feature_spec() # Because input schema includes unnecessary fields like 'species' and # 'island', we filter feature_spec to include required keys only. required_feature_spec = { k: v for k, v in feature_spec.items() if k in _FEATURE_KEYS } parsed_features = tf.io.parse_example(serialized_tf_examples, required_feature_spec) # Preprocess parsed input with transform operation defined in # preprocessing_fn(). transformed_features, _ = _apply_preprocessing(parsed_features, model.tft_layer) # Run inference with ML model. return model(transformed_features) return serve_tf_examples_fn def _input_fn(file_pattern: List[Text], data_accessor: tfx.components.DataAccessor, tf_transform_output: tft.TFTransformOutput, batch_size: int = 200) -> tf.data.Dataset: """Generates features and label for tuning/training. Args: file_pattern: List of paths or patterns of input tfrecord files. data_accessor: DataAccessor for converting input to RecordBatch. tf_transform_output: A TFTransformOutput. batch_size: representing the number of consecutive elements of returned dataset to combine in a single batch Returns: A dataset that contains (features, indices) tuple where features is a dictionary of Tensors, and indices is a single Tensor of label indices. """ dataset = data_accessor.tf_dataset_factory( file_pattern, tfxio.TensorFlowDatasetOptions(batch_size=batch_size), schema=tf_transform_output.raw_metadata.schema) transform_layer = tf_transform_output.transform_features_layer() def apply_transform(raw_features): return _apply_preprocessing(raw_features, transform_layer) return dataset.map(apply_transform).repeat() def _build_keras_model() -> tf.keras.Model: """Creates a DNN Keras model for classifying penguin data. Returns: A Keras Model. """ # The model below is built with Functional API, please refer to # https://www.tensorflow.org/guide/keras/overview for all API options. inputs = [ keras.layers.Input(shape=(1,), name=key) for key in _FEATURE_KEYS ] d = keras.layers.concatenate(inputs) for _ in range(2): d = keras.layers.Dense(8, activation='relu')(d) outputs = keras.layers.Dense(3)(d) model = keras.Model(inputs=inputs, outputs=outputs) model.compile( optimizer=keras.optimizers.Adam(1e-2), loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True), metrics=[keras.metrics.SparseCategoricalAccuracy()]) model.summary(print_fn=logging.info) return model # TFX Trainer will call this function. def run_fn(fn_args: tfx.components.FnArgs): """Train the model based on given args. Args: fn_args: Holds args used to train the model as name/value pairs. """ tf_transform_output = tft.TFTransformOutput(fn_args.transform_output) train_dataset = _input_fn( fn_args.train_files, fn_args.data_accessor, tf_transform_output, batch_size=_TRAIN_BATCH_SIZE) eval_dataset = _input_fn( fn_args.eval_files, fn_args.data_accessor, tf_transform_output, batch_size=_EVAL_BATCH_SIZE) model = _build_keras_model() model.fit( train_dataset, steps_per_epoch=fn_args.train_steps, validation_data=eval_dataset, validation_steps=fn_args.eval_steps) # NEW: Save a computation graph including transform layer. signatures = { 'serving_default': _get_serve_tf_examples_fn(model, tf_transform_output), } model.save(fn_args.serving_model_dir, save_format='tf', signatures=signatures)

Now you have completed all of the preparation steps to build a TFX pipeline.

Write a pipeline definition¶

We define a function to create a TFX pipeline. A Pipeline object represents a TFX pipeline, which can be run using one of the pipeline orchestration systems that TFX supports.

In [ ]:

def _create_pipeline(pipeline_name: str, pipeline_root: str, data_root: str,
                     schema_path: str, module_file: str, serving_model_dir: str,
                     metadata_path: str) -> tfx.dsl.Pipeline:
  """Implements the penguin pipeline with TFX."""
  # Brings data into the pipeline or otherwise joins/converts training data.
  example_gen = tfx.components.CsvExampleGen(input_base=data_root)

  # Computes statistics over data for visualization and example validation.
  statistics_gen = tfx.components.StatisticsGen(
      examples=example_gen.outputs['examples'])

  # Import the schema.
  schema_importer = tfx.dsl.Importer(
      source_uri=schema_path,
      artifact_type=tfx.types.standard_artifacts.Schema).with_id(
          'schema_importer')

  # Performs anomaly detection based on statistics and data schema.
  example_validator = tfx.components.ExampleValidator(
      statistics=statistics_gen.outputs['statistics'],
      schema=schema_importer.outputs['result'])

  # NEW: Transforms input data using preprocessing_fn in the 'module_file'.
  transform = tfx.components.Transform(
      examples=example_gen.outputs['examples'],
      schema=schema_importer.outputs['result'],
      materialize=False,
      module_file=module_file)

  # Uses user-provided Python function that trains a model.
  trainer = tfx.components.Trainer(
      module_file=module_file,
      examples=example_gen.outputs['examples'],

      # NEW: Pass transform_graph to the trainer.
      transform_graph=transform.outputs['transform_graph'],

      train_args=tfx.proto.TrainArgs(num_steps=100),
      eval_args=tfx.proto.EvalArgs(num_steps=5))

  # Pushes the model to a filesystem destination.
  pusher = tfx.components.Pusher(
      model=trainer.outputs['model'],
      push_destination=tfx.proto.PushDestination(
          filesystem=tfx.proto.PushDestination.Filesystem(
              base_directory=serving_model_dir)))

  components = [
      example_gen,
      statistics_gen,
      schema_importer,
      example_validator,

      transform,  # NEW: Transform component was added to the pipeline.

      trainer,
      pusher,
  ]

  return tfx.dsl.Pipeline(
      pipeline_name=pipeline_name,
      pipeline_root=pipeline_root,
      metadata_connection_config=tfx.orchestration.metadata
      .sqlite_metadata_connection_config(metadata_path),
      components=components)

def _create_pipeline(pipeline_name: str, pipeline_root: str, data_root: str, schema_path: str, module_file: str, serving_model_dir: str, metadata_path: str) -> tfx.dsl.Pipeline: """Implements the penguin pipeline with TFX.""" # Brings data into the pipeline or otherwise joins/converts training data. example_gen = tfx.components.CsvExampleGen(input_base=data_root) # Computes statistics over data for visualization and example validation. statistics_gen = tfx.components.StatisticsGen( examples=example_gen.outputs['examples']) # Import the schema. schema_importer = tfx.dsl.Importer( source_uri=schema_path, artifact_type=tfx.types.standard_artifacts.Schema).with_id( 'schema_importer') # Performs anomaly detection based on statistics and data schema. example_validator = tfx.components.ExampleValidator( statistics=statistics_gen.outputs['statistics'], schema=schema_importer.outputs['result']) # NEW: Transforms input data using preprocessing_fn in the 'module_file'. transform = tfx.components.Transform( examples=example_gen.outputs['examples'], schema=schema_importer.outputs['result'], materialize=False, module_file=module_file) # Uses user-provided Python function that trains a model. trainer = tfx.components.Trainer( module_file=module_file, examples=example_gen.outputs['examples'], # NEW: Pass transform_graph to the trainer. transform_graph=transform.outputs['transform_graph'], train_args=tfx.proto.TrainArgs(num_steps=100), eval_args=tfx.proto.EvalArgs(num_steps=5)) # Pushes the model to a filesystem destination. pusher = tfx.components.Pusher( model=trainer.outputs['model'], push_destination=tfx.proto.PushDestination( filesystem=tfx.proto.PushDestination.Filesystem( base_directory=serving_model_dir))) components = [ example_gen, statistics_gen, schema_importer, example_validator, transform, # NEW: Transform component was added to the pipeline. trainer, pusher, ] return tfx.dsl.Pipeline( pipeline_name=pipeline_name, pipeline_root=pipeline_root, metadata_connection_config=tfx.orchestration.metadata .sqlite_metadata_connection_config(metadata_path), components=components)

Run the pipeline¶

We will use LocalDagRunner as in the previous tutorial.

In [ ]:

tfx.orchestration.LocalDagRunner().run(
  _create_pipeline(
      pipeline_name=PIPELINE_NAME,
      pipeline_root=PIPELINE_ROOT,
      data_root=DATA_ROOT,
      schema_path=SCHEMA_PATH,
      module_file=_module_file,
      serving_model_dir=SERVING_MODEL_DIR,
      metadata_path=METADATA_PATH))

tfx.orchestration.LocalDagRunner().run( _create_pipeline( pipeline_name=PIPELINE_NAME, pipeline_root=PIPELINE_ROOT, data_root=DATA_ROOT, schema_path=SCHEMA_PATH, module_file=_module_file, serving_model_dir=SERVING_MODEL_DIR, metadata_path=METADATA_PATH))

You should see "INFO:absl:Component Pusher is finished." if the pipeline finished successfully.

The pusher component pushes the trained model to the SERVING_MODEL_DIR which is the serving_model/penguin-transform directory if you did not change the variables in the previous steps. You can see the result from the file browser in the left-side panel in Colab, or using the following command:

In [ ]:

# List files in created model directory.
!find {SERVING_MODEL_DIR}

# List files in created model directory. !find {SERVING_MODEL_DIR}

You can also check the signature of the generated model using the saved_model_cli tool.

In [ ]:

!saved_model_cli show --dir {SERVING_MODEL_DIR}/$(ls -1 {SERVING_MODEL_DIR} | sort -nr | head -1) --tag_set serve --signature_def serving_default

!saved_model_cli show --dir {SERVING_MODEL_DIR}/$(ls -1 {SERVING_MODEL_DIR} | sort -nr | head -1) --tag_set serve --signature_def serving_default

Because we defined serving_default with our own serve_tf_examples_fn function, the signature shows that it takes a single string. This string is a serialized string of tf.Examples and will be parsed with the tf.io.parse_example() function as we defined earlier (learn more about tf.Examples here).

We can load the exported model and try some inferences with a few examples.

In [ ]:

# Find a model with the latest timestamp.
model_dirs = (item for item in os.scandir(SERVING_MODEL_DIR) if item.is_dir())
model_path = max(model_dirs, key=lambda i: int(i.name)).path

loaded_model = tf.keras.models.load_model(model_path)
inference_fn = loaded_model.signatures['serving_default']

# Find a model with the latest timestamp. model_dirs = (item for item in os.scandir(SERVING_MODEL_DIR) if item.is_dir()) model_path = max(model_dirs, key=lambda i: int(i.name)).path loaded_model = tf.keras.models.load_model(model_path) inference_fn = loaded_model.signatures['serving_default']

In [ ]:

# Prepare an example and run inference.
features = {
  'culmen_length_mm': tf.train.Feature(float_list=tf.train.FloatList(value=[49.9])),
  'culmen_depth_mm': tf.train.Feature(float_list=tf.train.FloatList(value=[16.1])),
  'flipper_length_mm': tf.train.Feature(int64_list=tf.train.Int64List(value=[213])),
  'body_mass_g': tf.train.Feature(int64_list=tf.train.Int64List(value=[5400])),
}
example_proto = tf.train.Example(features=tf.train.Features(feature=features))
examples = example_proto.SerializeToString()

result = inference_fn(examples=tf.constant([examples]))
print(result['output_0'].numpy())

# Prepare an example and run inference. features = { 'culmen_length_mm': tf.train.Feature(float_list=tf.train.FloatList(value=[49.9])), 'culmen_depth_mm': tf.train.Feature(float_list=tf.train.FloatList(value=[16.1])), 'flipper_length_mm': tf.train.Feature(int64_list=tf.train.Int64List(value=[213])), 'body_mass_g': tf.train.Feature(int64_list=tf.train.Int64List(value=[5400])), } example_proto = tf.train.Example(features=tf.train.Features(feature=features)) examples = example_proto.SerializeToString() result = inference_fn(examples=tf.constant([examples])) print(result['output_0'].numpy())

The third element, which corresponds to 'Gentoo' species, is expected to be the largest among three.

Next steps¶

If you want to learn more about Transform component, see Transform Component guide. You can find more resources on https://www.tensorflow.org/tfx/tutorials.

Please see Understanding TFX Pipelines to learn more about various concepts in TFX.