TFX の TensorFlow 2.x

TensorFlow 2.0 は 2019 年にリリースされ、 Keras の緊密な統合、デフォルトでの積極的な実行、およびPythonic 関数の実行、その他の新機能と改善が組み込まれました。

このガイドでは、TFX の TF 2.x の包括的な技術概要を説明します。

どのバージョンを使用するか?

TFX は TensorFlow 2.x と互換性があり、TensorFlow 1.x に存在していた高レベル API (特に Estimator) は引き続き動作します。

TensorFlow 2.x で新しいプロジェクトを開始する

TensorFlow 2.x は TensorFlow 1.x の高レベルの機能を保持しているため、新しい機能を使用する予定がない場合でも、新しいプロジェクトで古いバージョンを使用する利点はありません。

したがって、新しい TFX プロジェクトを開始する場合は、TensorFlow 2.x を使用することをお勧めします。 Keras やその他の新機能の完全なサポートが利用可能になったら、後でコードを更新することもできます。また、TensorFlow 1.x からアップグレードしようとするのではなく、TensorFlow 2.x から始めた場合、変更の範囲ははるかに制限されます。未来。

既存のプロジェクトを TensorFlow 2.x に変換する

TensorFlow 1.x 用に書かれたコードは TensorFlow 2.x とほぼ互換性があり、TFX でも引き続き動作します。

ただし、TF 2.x で利用可能になった改良点や新機能を活用したい場合は、 TF 2.x への移行手順に従うことができます。

推定者

Estimator API は TensorFlow 2.x でも維持されていますが、新しい機能や開発の焦点では​​ありません。 Estimator を使用して TensorFlow 1.x または 2.x で書かれたコードは、引き続き TFX で期待どおりに動作します。

以下は、純粋な Estimator を使用したエンドツーエンドの TFX の例です:タクシーの例 (Estimator)

Keras とmodel_to_estimator

Keras モデルはtf.keras.estimator.model_to_estimator関数でラップすることができ、これにより Estimator であるかのように動作できるようになります。これを使用するには:

  1. Keras モデルを構築します。
  2. コンパイルされたモデルをmodel_to_estimatorに渡します。
  3. 通常の Estimator の使用方法と同じように、 model_to_estimatorの結果を Trainer で使用します。 
# Build a Keras model.
def _keras_model_builder():
  """Creates a Keras model."""
  ...

  model = tf.keras.Model(inputs=inputs, outputs=output)
  model.compile()

  return model


# Write a typical trainer function
def trainer_fn(trainer_fn_args, schema):
  """Build the estimator, using model_to_estimator."""
  ...

  # Model to estimator
  estimator = tf.keras.estimator.model_to_estimator(
      keras_model=_keras_model_builder(), config=run_config)

  return {
      'estimator': estimator,
      ...
  }

Trainer のユーザー モジュール ファイルを除き、パイプラインの残りの部分は変更されません。

ネイティブ Keras (つまり、 model_to_estimatorのない Keras)

例とコラボ

ネイティブ Keras を使用した例をいくつか示します。

コンポーネントごとのKeras Colabもあります。

TFX コンポーネント

次のセクションでは、関連する TFX コンポーネントがネイティブ Keras をサポートする方法について説明します。

変身

Transform は現在、Keras モデルを実験的にサポートしています。

Transform コンポーネント自体は、変更せずにネイティブ Keras で使用できます。 preprocessing_fn定義は同じままで、 TensorFlowおよびtf.Transform演算を使用します。

ネイティブ Keras 用に、serving 関数と eval 関数が変更されました。詳細については、次のトレーナーと評価者のセクションで説明します。

トレーナー

ネイティブ Keras を構成するには、 GenericExecutor Trainer コンポーネントに設定して、デフォルトの Estimator ベースの Executor を置き換える必要があります。詳細はこちらをご確認ください。

Transform を含む Keras モジュール ファイル

トレーニング モジュール ファイルには、 GenericExecutorによって呼び出されるrun_fn含まれている必要があります。典型的な Keras run_fn次のようになります。

def run_fn(fn_args: TrainerFnArgs):
  """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 and eval files contains transformed examples.
  # _input_fn read dataset based on transformed schema from tft.
  train_dataset = _input_fn(fn_args.train_files, fn_args.data_accessor,
                            tf_transform_output.transformed_metadata.schema)
  eval_dataset = _input_fn(fn_args.eval_files, fn_args.data_accessor,
                           tf_transform_output.transformed_metadata.schema)

  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)

  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)

上記のrun_fnでは、モデルが予測用に生のサンプルを取得できるように、トレーニング済みモデルをエクスポートするときにサービス署名が必要です。典型的なサービス提供関数は次のようになります。

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

  # the layer is added as an attribute to the model in order to make sure that
  # the model assets are handled correctly when exporting.
  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

上記の提供関数では、 tft.TransformFeaturesLayerレイヤーを使用して、推論用の生データに tf.Transform 変換を適用する必要があります。 Estimator に必要だった以前の_serving_input_receiver_fnは、Keras では必要なくなります。

変換なしの Keras モジュール ファイル

これは上記のモジュール ファイルに似ていますが、変換は行われていません。

def _get_serve_tf_examples_fn(model, schema):

  @tf.function
  def serve_tf_examples_fn(serialized_tf_examples):
    feature_spec = _get_raw_feature_spec(schema)
    feature_spec.pop(_LABEL_KEY)
    parsed_features = tf.io.parse_example(serialized_tf_examples, feature_spec)
    return model(parsed_features)

  return serve_tf_examples_fn


def run_fn(fn_args: TrainerFnArgs):
  schema = io_utils.parse_pbtxt_file(fn_args.schema_file, schema_pb2.Schema())

  # Train and eval files contains raw examples.
  # _input_fn reads the dataset based on raw data schema.
  train_dataset = _input_fn(fn_args.train_files, fn_args.data_accessor, schema)
  eval_dataset = _input_fn(fn_args.eval_files, fn_args.data_accessor, schema)

  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)

  signatures = {
      'serving_default':
          _get_serve_tf_examples_fn(model, 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)
tf.distribute.Strategy

現時点では、TFX は単一ワーカー戦略 (例: MirroredStrategy 、 OneDeviceStrategy ) のみをサポートしています。

配布戦略を使用するには、適切な tf.distribute.Strategy を作成し、Keras モデルの作成とコンパイルを戦略スコープ内に移動します。

たとえば、上記のmodel = _build_keras_model()を次のように置き換えます。

  mirrored_strategy = tf.distribute.MirroredStrategy()
  with mirrored_strategy.scope():
    model = _build_keras_model()

  # Rest of the code can be unchanged.
  model.fit(...)

MirroredStrategyによって使用されるデバイス (CPU/GPU) を確認するには、情報レベルの tensorflow ログを有効にします。

import logging
logging.getLogger("tensorflow").setLevel(logging.INFO)

ログにUsing MirroredStrategy with devices (...)表示されるはずです。

評価者

TFMA v0.2x では、ModelValidator と Evaluator が 1 つの新しい Evaluator コンポーネントに統合されました。新しい Evaluator コンポーネントは、単一モデルの評価を実行できるほか、現在のモデルを以前のモデルと比較して検証することもできます。この変更により、Pusher コンポーネントは ModelValidator ではなく Evaluator からの祝福結果を消費するようになりました。

新しい Evaluator は、Estimator モデルだけでなく Keras モデルもサポートしています。 Evaluator は提供に使用されるのと同じSavedModelに基づいているため、以前は必要であった_eval_input_receiver_fnおよび eval 保存モデルは Keras では必要なくなりました。

詳細については、「評価者」を参照してください。