Day 29 : 用於生产的 TensorFlow Extended (TFX) 实作

用於生产的机械学习系统，在 Day 28 介绍 TensorFlow Extended (TFX) 解决方案，是专门用於可扩充的高效能机器学习工作，包括建立模型、进行训练、提供推论，以及管理线上、行动装置 (TensorFlow Lite)和网页应用服务 (TensorFlow JS) 的部署。
本日我们修改 TFX 官方在笔记本可执行的互动式范例，透过实作理解 TFX 如何在工作流程中透过组件功能验证资料、特徵工程、训练模型、模型分析到部署模型。
Colab 实作范例，内容源自TFX 官方范例。

TFX 组件笔记本互动实作

0. 安装与设置 TFX 环境

更新 Colab 的 pip。
安装 TFX，安装完需重新启动执行阶段(Restart Runtime)
```
!pip install tfx
```
设定工作路径。
下载芝加哥 Taxi Trips 资料集，将集构建一个预测小费tips的模型。

0. 创建互动文件检视器 InteractiveContext

tfx.orchestration.experimental.interactive.interactive_context.InteractiveContext 允许在 notebook 环境中以互动方式查看 TFX 组件。
InteractiveContext 预设使用临时的中继资料。
- 已有自己的 pipeline 可设定 pipe_root 参数。
- 已有中继资料库可设定 metadata_connection_config 参数。
透过 InteractiveContext() 逐一演示互动情形，请留意 InteractiveContext.run() 是在笔记本互动显示方式，在生产环境中可专注组件流程（请参阅构建 TFX 管道指南）。

1. ExampleGen

ExampleGen 将数据拆分为训练集和评估集（预设为 2/3 训练、 1/3 评估）
ExampleGen 将数据转换为 tf.Example 格式（参阅说明）。

本范例将 _data_root 的 CSV 资料集输入至 ExampleGen。

example_gen = tfx.components.CsvExampleGen(input_base=_data_root)
context.run(example_gen)

观察互动介面内容，建议您使用 Colab 测试，在此介面中:
- .execution_id 是持续累加的版次编号，在对应的资料夹将各版次的内容留存纪录。
- .component 指的是该 TFX 组件，譬如下图显示为 .component.CsvExampleGen。
- .component.inputs 纪录输入来源。
- .component.outputs 纪录输出结果。

2. StatisticsGen

StatisticsGen 组件输入 ExampleGen 数据後，将据以计算出资料集的统计数据。

StatisticsGen 是 TFDV 模组功能之一。

statistics_gen = tfx.components.StatisticsGen(
    examples = example_gen.outputs['examples']
    )
context.run(statistics_gen)

context.run(statistics_gen) 观察互动介面:
- .execution_id 版次累加至2。
- .component.inputs 组件输入为 Examples 。
- 输出为 ExampleStatistics 。
context.show(statistics_gen.outputs['statistics']) 如同 TFDV 工具以 Facets 视觉化统计资讯。
```
context.show(statistics_gen.outputs['statistics'])
```
- 可以观察判读可能异常的红色值、资料分布情形等。

3. SchemaGen

SchemaGen组件会依据您的资料统计自动产生 Schema ，包含数据预期边界、资料类型与属性它还使用TensorFlow 数据验证库。
SchemaGen 同样是 TFDV 模组功能之一。
即便 Schema 自动生成已经很实用，但您仍应该会依据需求进行审查和修改。
```
schema_gen = tfx.components.SchemaGen(
    statistics=statistics_gen.outputs['statistics'],
    infer_feature_shape=False)
context.run(schema_gen)
```
- SchemaGen 输入为 StatisticsGen，默认情况下查看已拆分的训练资料集。
- SchemaGen 输出为 Schema 。
- SchemaGen 执行後可透过 context.show(schema_gen.outputs['schema']) 查看 Schema 表格。
- 表格呈现各特徵名称、属性、是否必须、所有值、Domain 及边界范围等，参见 SchemaGen 文件。

4. ExampleValidator

ExampleValidator 组件根据 Schema 的预期检测数据中的异常。
ExampleValidator 同样是 TFDV 模组功能之一。
ExampleValidator 的输入是来自具有数据统计资讯的 StatisticsGen 以及具有数据定义 Schema 的 SchemaGen。

ExampleValidator 的输出 anomalies 是有无异常的判读结果。

example_validator = tfx.components.ExampleValidator(
    statistics = statistics_gen.outputs['statistics'],
    schema = schema_gen.outputs['schema'])
context.run(example_validator)

执行 ExampleValidator 後可以产生异常情形的图表，绿字 No anomalies found. 表示无异常。
```
context.show(example_validator.outputs['anomalies'])
```

5. Transform

Transform 组件为训练和服务执行特徵工程。
Transform 使用TensorFlow Transform 模组。
Transform 输入数据来自 ExampleGen 、 Schema 来自 SchemaGen ，以及自行定义如何进行特徵工程的模组。
以下为自行定义的 Transform 程序码范例，（有关 TensorFlow Transform API 的介绍，请参阅教程）。

Notebook 魔术指令 %%writefile ，可以将 cell 内的程序码指定保存为档案，该档案可以用 Transform 组件将程序码档案做为模组输入执行。

%%writefile {_taxi_constants_module_file}

# 假设分类特徵的最大值
MAX_CATEGORICAL_FEATURE_VALUES = [24, 31, 12]

CATEGORICAL_FEATURE_KEYS = [
    'trip_start_hour', 
    'trip_start_day', 
    'trip_start_month',
    'pickup_census_tract', 
    'dropoff_census_tract', 
    'pickup_community_area',
    'dropoff_community_area'
    ]

DENSE_FLOAT_FEATURE_KEYS = ['trip_miles', 'fare', 'trip_seconds']

# tf.transform用於编码每个特徵的桶数=10
FEATURE_BUCKET_COUNT = 10

BUCKET_FEATURE_KEYS = [
    'pickup_latitude', 
    'pickup_longitude', 
    'dropoff_latitude',
    'dropoff_longitude'
    ]

# tf.transform用於编码VOCAB_FEATURES的词汇术语数量=1000
VOCAB_SIZE = 1000

# Count of out-of-vocab buckets in which unrecognized 
# VOCAB_FEATURES are hashed.
OOV_SIZE = 10

VOCAB_FEATURE_KEYS = [
    'payment_type',
    'company',
]

# Keys
LABEL_KEY = 'tips'
FARE_KEY = 'fare'

def transformed_name(key):
  return key + '_xf'

接着编写 preprocessing_fn 将原始数据转换特徵。

%%writefile {_taxi_transform_module_file}

import tensorflow as tf
import tensorflow_transform as tft

import taxi_constants

_DENSE_FLOAT_FEATURE_KEYS = taxi_constants.DENSE_FLOAT_FEATURE_KEYS
_VOCAB_FEATURE_KEYS = taxi_constants.VOCAB_FEATURE_KEYS
_VOCAB_SIZE = taxi_constants.VOCAB_SIZE
_OOV_SIZE = taxi_constants.OOV_SIZE
_FEATURE_BUCKET_COUNT = taxi_constants.FEATURE_BUCKET_COUNT
_BUCKET_FEATURE_KEYS = taxi_constants.BUCKET_FEATURE_KEYS
_CATEGORICAL_FEATURE_KEYS = taxi_constants.CATEGORICAL_FEATURE_KEYS
_FARE_KEY = taxi_constants.FARE_KEY
_LABEL_KEY = taxi_constants.LABEL_KEY
_transformed_name = taxi_constants.transformed_name


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 
    operations.
  """
  outputs = {}
  for key in _DENSE_FLOAT_FEATURE_KEYS:
    # Preserve this feature as a dense float, 
    # setting nan's to the mean.
    outputs[_transformed_name(key)] = tft.scale_to_z_score(
        _fill_in_missing(inputs[key]))

  for key in _VOCAB_FEATURE_KEYS:
    # Build a vocabulary for this feature.
    outputs[_transformed_name(key)] = tft.compute_and_apply_vocabulary(
        _fill_in_missing(inputs[key]),
        top_k=_VOCAB_SIZE,
        num_oov_buckets=_OOV_SIZE)

  for key in _BUCKET_FEATURE_KEYS:
    outputs[_transformed_name(key)] = tft.bucketize(
        _fill_in_missing(inputs[key]), _FEATURE_BUCKET_COUNT)

  for key in _CATEGORICAL_FEATURE_KEYS:
    outputs[_transformed_name(key)] = _fill_in_missing(inputs[key])

  # Was this passenger a big tipper?
  taxi_fare = _fill_in_missing(inputs[_FARE_KEY])
  tips = _fill_in_missing(inputs[_LABEL_KEY])
  outputs[_transformed_name(_LABEL_KEY)] = tf.where(
      tf.math.is_nan(taxi_fare),
      tf.cast(
          tf.zeros_like(taxi_fare), 
          tf.int64
          ),
      # Test if the tip was > 20% of the fare.
      tf.cast(
          tf.greater(
          tips, 
          tf.multiply(taxi_fare, tf.constant(0.2))), 
          tf.int64
          )
      )

  return outputs


def _fill_in_missing(x):
  """Replace missing values in a SparseTensor.
  Fills in missing values of `x` with '' or 0, 
  and converts to a dense tensor.

  Args:
    x: A `SparseTensor` of rank 2.  
    Its dense shape should have size at most 1 in the second dimension.
  Returns:
    A rank 1 tensor where missing values of `x` have been filled in.
  """
  if not isinstance(x, tf.sparse.SparseTensor):
    return x

  default_value = '' if x.dtype == tf.string else 0
  return tf.squeeze(
      tf.sparse.to_dense(
          tf.SparseTensor(x.indices, x.values, [x.dense_shape[0], 1]),
          default_value),
          axis=1
          )

将特徵工程程序传递给 Transform 组件转换资料。

transform = tfx.components.Transform(
    examples=example_gen.outputs['examples'],
    schema=schema_gen.outputs['schema'],
    module_file=os.path.abspath(_taxi_transform_module_file))
context.run(transform)

Transform 组件将产生以下两种类型的输出：
- transform_graph 是可以执行预处理操作的图（此图将包含在服务和评估模型中）。
- transformed_examples 表示预处理的训练和评估数据。
- 查看transform.outputs 。
输出的 transform_graph 同时指向包含3个子目录的目录。
- transformed_metadata子目录包含预处理数据的架构。
- transform_fn子目录包含实际的预处理图。
- metadata子目录包含原始数据的架构。
```
train_uri = transform.outputs['transform_graph'].get()[0].uri
os.listdir(train_uri)
```

6. Trainer

Trainer组件负责训练 TensorFlow 模型。
Trainer 预设使用 Estimator API ，如要使用 Keras API，您需要通过在 Trainer 的构造函数中设置来指定 custom_executor_spec=executor_spec.ExecutorClassSpec(GenericExecutor) ，参阅Generic Trainer 。
Trainer 的输入来源:
- 来自 SchemaGen 的 Schema。
- 来自 Transform 的 graph。
- 训练参数。
- 做为模组输入的自定义程序码。

以下为用户自定义模型代码示范（参见 TensorFlow Keras API 介绍）。

%%writefile {_taxi_trainer_module_file}

from typing import List, Text

import os
import absl
import datetime
import tensorflow as tf
import tensorflow_transform as tft

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

import taxi_constants

_DENSE_FLOAT_FEATURE_KEYS = taxi_constants.DENSE_FLOAT_FEATURE_KEYS
_VOCAB_FEATURE_KEYS = taxi_constants.VOCAB_FEATURE_KEYS
_VOCAB_SIZE = taxi_constants.VOCAB_SIZE
_OOV_SIZE = taxi_constants.OOV_SIZE
_FEATURE_BUCKET_COUNT = taxi_constants.FEATURE_BUCKET_COUNT
_BUCKET_FEATURE_KEYS = taxi_constants.BUCKET_FEATURE_KEYS
_CATEGORICAL_FEATURE_KEYS = taxi_constants.CATEGORICAL_FEATURE_KEYS
_MAX_CATEGORICAL_FEATURE_VALUES = taxi_constants.MAX_CATEGORICAL_FEATURE_VALUES
_LABEL_KEY = taxi_constants.LABEL_KEY
_transformed_name = taxi_constants.transformed_name


def _transformed_names(keys):
  return [_transformed_name(key) for key in keys]


def _get_serve_tf_examples_fn(model, tf_transform_output):
  """Returns a function that parses a serialized tf.Example and applies TFT."""

  model.tft_layer = 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)
    parsed_features = tf.io.parse_example(serialized_tf_examples, feature_spec)
    transformed_features = model.tft_layer(parsed_features)
    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.
  """
  return data_accessor.tf_dataset_factory(
      file_pattern,
      tfxio.TensorFlowDatasetOptions(
          batch_size=batch_size, label_key=_transformed_name(_LABEL_KEY)),
      tf_transform_output.transformed_metadata.schema)


def _build_keras_model(hidden_units: List[int] = None) -> tf.keras.Model:
  """Creates a DNN Keras model for classifying taxi data.
  Args:
    hidden_units: [int], the layer sizes of the DNN (input layer first).
  Returns:
    A keras Model.
  """
  real_valued_columns = [
      tf.feature_column.numeric_column(key, shape=())
      for key in _transformed_names(_DENSE_FLOAT_FEATURE_KEYS)
  ]
  categorical_columns = [
      tf.feature_column.categorical_column_with_identity(
          key, num_buckets=_VOCAB_SIZE + _OOV_SIZE, default_value=0)
      for key in _transformed_names(_VOCAB_FEATURE_KEYS)
  ]
  categorical_columns += [
      tf.feature_column.categorical_column_with_identity(
          key, num_buckets=_FEATURE_BUCKET_COUNT, default_value=0)
      for key in _transformed_names(_BUCKET_FEATURE_KEYS)
  ]
  categorical_columns += [
      tf.feature_column.categorical_column_with_identity(  # pylint: disable=g-complex-comprehension
          key,
          num_buckets=num_buckets,
          default_value=0) for key, num_buckets in zip(
              _transformed_names(_CATEGORICAL_FEATURE_KEYS),
              _MAX_CATEGORICAL_FEATURE_VALUES)
  ]
  indicator_column = [
      tf.feature_column.indicator_column(categorical_column)
      for categorical_column in categorical_columns
  ]

  model = _wide_and_deep_classifier(
      # TODO(b/139668410) replace with premade wide_and_deep keras model
      wide_columns=indicator_column,
      deep_columns=real_valued_columns,
      dnn_hidden_units=hidden_units or [100, 70, 50, 25])
  return model


def _wide_and_deep_classifier(wide_columns, deep_columns, dnn_hidden_units):
  """Build a simple keras wide and deep model.

  Args:
    wide_columns: Feature columns wrapped in indicator_column for wide (linear)
      part of the model.
    deep_columns: Feature columns for deep part of the model.
    dnn_hidden_units: [int], the layer sizes of the hidden DNN.

  Returns:
    A Wide and Deep Keras model
  """
  # Following values are hard coded for simplicity in this example,
  # However prefarably they should be passsed in as hparams.

  # Keras needs the feature definitions at compile time.
  # TODO(b/139081439): Automate generation of input layers from FeatureColumn.
  input_layers = {
      colname: tf.keras.layers.Input(
          name=colname, 
          shape=(), 
          dtype=tf.float32
          )
      for colname in _transformed_names(_DENSE_FLOAT_FEATURE_KEYS)
  }
  input_layers.update({
      colname: tf.keras.layers.Input(
          name=colname, 
          shape=(), 
          dtype='int32')
      for colname in _transformed_names(_VOCAB_FEATURE_KEYS)
  })
  input_layers.update({
      colname: tf.keras.layers.Input(
          name=colname, 
          shape=(), 
          dtype='int32')
      for colname in _transformed_names(_BUCKET_FEATURE_KEYS)
  })
  input_layers.update({
      colname: tf.keras.layers.Input(
          name=colname, 
          shape=(), 
          dtype='int32')
      for colname in _transformed_names(_CATEGORICAL_FEATURE_KEYS)
  })

  # TODO(b/161952382): Replace with Keras preprocessing layers.
  deep = tf.keras.layers.DenseFeatures(deep_columns)(input_layers)
  for numnodes in dnn_hidden_units:
    deep = tf.keras.layers.Dense(numnodes)(deep)
  wide = tf.keras.layers.DenseFeatures(wide_columns)(input_layers)

  output = tf.keras.layers.Dense(1)(
          tf.keras.layers.concatenate([deep, wide]))

  model = tf.keras.Model(input_layers, output)
  model.compile(
      loss=tf.keras.losses.BinaryCrossentropy(from_logits=True),
      optimizer=tf.keras.optimizers.Adam(lr=0.001),
      metrics=[tf.keras.metrics.BinaryAccuracy()])
  model.summary(print_fn=absl.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.
  """
  # Number of nodes in the first layer of the DNN
  first_dnn_layer_size = 100
  num_dnn_layers = 4
  dnn_decay_factor = 0.7

  tf_transform_output = tft.TFTransformOutput(fn_args.transform_output)

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

  model = _build_keras_model(
      # Construct layers sizes with exponetial decay
      hidden_units=[
          max(2, int(first_dnn_layer_size * dnn_decay_factor**i))
          for i in range(num_dnn_layers)
      ])

  tensorboard_callback = tf.keras.callbacks.TensorBoard(
      log_dir=fn_args.model_run_dir, update_freq='batch')
  model.fit(
      train_dataset,
      steps_per_epoch=fn_args.train_steps,
      validation_data=eval_dataset,
      validation_steps=fn_args.eval_steps,
      callbacks=[tensorboard_callback])

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

创立 taxi_trainer.py 之後将程序码做为模组传递给 Trainer 组件并运行它来训练模型。

trainer = tfx.components.Trainer(
    module_file=os.path.abspath(_taxi_trainer_module_file),
    examples=transform.outputs['transformed_examples'],
    transform_graph=transform.outputs['transform_graph'],
    schema=schema_gen.outputs['schema'],
    train_args=tfx.proto.TrainArgs(num_steps=10000),
    eval_args=tfx.proto.EvalArgs(num_steps=5000))
context.run(trainer)

互动内文如下:

(选用) 以 TensorBoard 分析训练模型

可以透过 TensorBoard 分析模型训练曲线。

model_run_artifact_dir = trainer.outputs['model_run'].get()[0].uri

%load_ext tensorboard
%tensorboard --logdir {model_run_artifact_dir}

视觉化结果如下:

7. Evaluator

Evaluator 组件可评估模型性能。
Evaluator 组件为 TensorFlow Model Analysis (TFMA) 模组功能。

Evaluator 可以设定门槛值以比较并选择较佳的模型。这在生产管道设置中很有用，您可以每天自动训练和验证模型。

model_resolver = tfx.dsl.Resolver(
      strategy_class=tfx.dsl.experimental.LatestBlessedModelStrategy,
      model=tfx.dsl.Channel(
          type=tfx.types.standard_artifacts.Model),
      model_blessing=tfx.dsl.Channel(
          type=tfx.types.standard_artifacts.ModelBlessing)).with_id(
              'latest_blessed_model_resolver')
context.run(model_resolver)

evaluator = tfx.components.Evaluator(
    examples=example_gen.outputs['examples'],
    model=trainer.outputs['model'],
    baseline_model=model_resolver.outputs['model'],
    eval_config=eval_config)
context.run(evaluator)

Evaluator 的输入:
- 输入资料集来自 ExampleGen。
- 训练模型来自 Trainer 和切片配置。切片配置允许您根据特徵值对指标进行切片（例如，您的模型在早上 8 点和晚上 8 点开始的出租车行程中表现如何？）。
在此笔记本范例只训练一个模型，所以Evaluator自动将模型标记为“Good”。
```
context.show(evaluator.outputs['evaluation'])
```
要切片显示模型情形，需使用 TFMA 模组。

在此示范将trip_start_hour切片视觉化，TFMA 支援许多其他可视化，例如公平指标和绘制模型性能的时间序列。要了解更多信息，请参阅教学。

import tensorflow_model_analysis as tfma

# Get the TFMA output result path and load the result.
PATH_TO_RESULT = evaluator.outputs['evaluation'].get()[0].uri
tfma_result = tfma.load_eval_result(PATH_TO_RESULT)

# Show data sliced along feature column trip_start_hour.
tfma.view.render_slicing_metrics(
    tfma_result, slicing_column='trip_start_hour')

切片示范如下:

通过门槛值的模型会得到祝福 blessing ，第一次预设会自动取得，之後持续训练过程会将取得祝福的模型再上线。
```
blessing_uri = evaluator.outputs['blessing'].get()[0].uri
!ls -l {blessing_uri}
```

8. Pusher

Pusher 组件通常位於 TFX 管道末端。
Pusher 组件检查模型是否已通过验证，如果是，则将模型导出至 _serving_model_dir。

Pusher 将以 SavedModel 格式导出您的模型。

pusher = tfx.components.Pusher(
    model=trainer.outputs['model'],
    model_blessing=evaluator.outputs['blessing'],
    push_destination=tfx.proto.PushDestination(
        filesystem=tfx.proto.PushDestination.Filesystem(
            base_directory=_serving_model_dir)))
context.run(pusher)

终於完成 TFX 所有组件的示范!
如果在 GCP 在 AI Platform 以 Kubeflow 部署 TFX，Pipeline 进一步实作可参阅 TFX on Google Cloud AI Platform Pipelines 视觉化如下方便查阅。

小结

在笔记本逐步互动式的完成 TFX 各组件的作业流程，实际上工作可以不用那麽复杂，在 GCP 的解决方案，可以用 Google Cloud AI Platform Pipeline ，是 Google Cloud AI Platform + TFX 容器化部署的结合，您可以参阅TFX: Production ML with TensorFlow in 2020 (TF Dev Summit '20) 影片说明，理解 TFX 各组件串起的用於生产的机械学习工作流程。