GitHub Repository: tensorflow/docs-l10n
Path: blob/master/site/ja/guide/migrate/tpu_embedding.ipynb
²⁵¹¹⁸ views

Kernel: Python 3

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.

TPU embedding_columns から TPUEmbedding レイヤーに移行する

このガイドでは、TPU での埋め込みトレーニングを、TPUEstimator を使用した TensorFlow 1 の embedding_column API から TensorFlow 2 の TPUStrategy を使用した TPUEmbedding レイヤー API に移行する方法を示します。

埋め込みは（大きな）行列です。これらは、疎な特徴空間から密なベクトルにマップするルックアップテーブルです。埋め込みは、効率的で密な表現を提供し、特徴間の複雑な類似性と関係を捉えます。

TensorFlow には、TPU でのトレーニング埋め込みの特別なサポートが含まれています。この TPU 固有の埋め込みサポートにより、単一の TPU デバイスのメモリよりも大きい埋め込みをトレーニングし、TPU で疎で不規則な入力を使用できます。

TensorFlow 1 では、tf.compat.v1.estimator.tpu.TPUEstimator は、TPU で提供するためのトレーニング、評価、予測、およびエクスポートをカプセル化する高レベル API です。tf.compat.v1.tpu.experimental.embedding_column を特にサポートしています。
トレーニングと評価には、モデル構築（tf.keras.Model）、オプティマイザ（tf.keras.optimizers.Optimizer）、および Model.fit または tf.function や tf.GradientTape を使用したカスタムトレーニングループなど、Keras API と互換性のある TPU 分散ストラテジー tf.distribute.TPUStrategy を使用します。

詳細については、tfrs.layers.embedding.TPUEmbedding レイヤーの API ドキュメント、tf.tpu.experimental.embedding.TableConfig を参照してください。また、追加情報については、tf .tpu.experimental.embedding.FeatureConfig ドキュメントを参照してください。tf.distribute.TPUStrategy の概要については、分散トレーニングガイドと TPU の使用ガイドを参照してください。 TPUEstimator から TPUStrategy に移行する場合は、TPU 移行ガイドを参照してください。

セットアップ

TensorFlow Recommenders をインストールし、必要なパッケージをインポートします。

In [ ]:

!pip install tensorflow-recommenders

In [ ]:

import tensorflow as tf
import tensorflow.compat.v1 as tf1

# TPUEmbedding layer is not part of TensorFlow.
import tensorflow_recommenders as tfrs

そして、デモ用に単純なデータセットを準備します。

In [ ]:

features = [[1., 1.5]]
embedding_features_indices = [[0, 0], [0, 1]]
embedding_features_values = [0, 5]
labels = [[0.3]]
eval_features = [[4., 4.5]]
eval_embedding_features_indices = [[0, 0], [0, 1]]
eval_embedding_features_values = [4, 3]
eval_labels = [[0.8]]

TensorFlow 1: TPUEstimator を使用して TPU の埋め込みをトレーニングする

TensorFlow 1 では、tf.compat.v1.tpu.experimental.embedding_column API を使用して TPU の埋め込みを設定し、tf.compat.v1.estimator.tpu.TPUEstimator を使用して TPU でモデルをトレーニング/評価します。

入力は、ゼロから TPU 埋め込みテーブルの語彙サイズまでの範囲の整数です。tf.feature_column.categorical_column_with_identity を使用して、入力をカテゴリカル ID にエンコードすることから始めます。入力特徴量は整数値であるため、key パラメータには "sparse_feature" を使用します。num_buckets は埋め込み表の語彙サイズ（10）です。

In [ ]:

embedding_id_column = (
      tf1.feature_column.categorical_column_with_identity(
          key="sparse_feature", num_buckets=10))

次に、疎なカテゴリカル入力を tpu.experimental.embedding_column を使用して密な表現に変換します。ここで、dimension は埋め込み表の幅です。num_buckets ごとに埋め込みベクトルを格納します。

In [ ]:

embedding_column = tf1.tpu.experimental.embedding_column(
    embedding_id_column, dimension=5)

ここで、tf.estimator.tpu.experimental.EmbeddingConfigSpec を介して TPU 固有の埋め込み構成を定義します。後で embedding_config_spec パラメータとして tf.estimator.tpu.TPUEstimator に渡します。

In [ ]:

embedding_config_spec = tf1.estimator.tpu.experimental.EmbeddingConfigSpec(
    feature_columns=(embedding_column,),
    optimization_parameters=(
        tf1.tpu.experimental.AdagradParameters(0.05)))

次に、TPUEstimator を使用するには、以下を定義します。

トレーニングデータの入力関数
評価データの評価入力関数
TPUEstimator に特徴とラベルを使用してトレーニング演算を定義する方法を指示するためのモデル関数

In [ ]:

def _input_fn(params):
  dataset = tf1.data.Dataset.from_tensor_slices((
      {"dense_feature": features,
       "sparse_feature": tf1.SparseTensor(
           embedding_features_indices,
           embedding_features_values, [1, 2])},
           labels))
  dataset = dataset.repeat()
  return dataset.batch(params['batch_size'], drop_remainder=True)

def _eval_input_fn(params):
  dataset = tf1.data.Dataset.from_tensor_slices((
      {"dense_feature": eval_features,
       "sparse_feature": tf1.SparseTensor(
           eval_embedding_features_indices,
           eval_embedding_features_values, [1, 2])},
           eval_labels))
  dataset = dataset.repeat()
  return dataset.batch(params['batch_size'], drop_remainder=True)

def _model_fn(features, labels, mode, params):
  embedding_features = tf1.keras.layers.DenseFeatures(embedding_column)(features)
  concatenated_features = tf1.keras.layers.Concatenate(axis=1)(
      [embedding_features, features["dense_feature"]])
  logits = tf1.layers.Dense(1)(concatenated_features)
  loss = tf1.losses.mean_squared_error(labels=labels, predictions=logits)
  optimizer = tf1.train.AdagradOptimizer(0.05)
  optimizer = tf1.tpu.CrossShardOptimizer(optimizer)
  train_op = optimizer.minimize(loss, global_step=tf1.train.get_global_step())
  return tf1.estimator.tpu.TPUEstimatorSpec(mode, loss=loss, train_op=train_op)

これらの関数を定義したら、クラスタ情報を提供する tf.distribute.cluster_resolver.TPUClusterResolver と tf.compat.v1.estimator.tpu.RunConfig オブジェクトを作成します。

定義したモデル関数に加えて、TPUEstimator を作成できるようになりました。ここでは、チェックポイントの保存をスキップしてフローを簡素化します。次に、TPUEstimator のトレーニングと評価の両方のバッチサイズを指定します。

In [ ]:

cluster_resolver = tf1.distribute.cluster_resolver.TPUClusterResolver(tpu='')
print("All devices: ", tf1.config.list_logical_devices('TPU'))

In [ ]:

tpu_config = tf1.estimator.tpu.TPUConfig(
    iterations_per_loop=10,
    per_host_input_for_training=tf1.estimator.tpu.InputPipelineConfig
          .PER_HOST_V2)
config = tf1.estimator.tpu.RunConfig(
    cluster=cluster_resolver,
    save_checkpoints_steps=None,
    tpu_config=tpu_config)
estimator = tf1.estimator.tpu.TPUEstimator(
    model_fn=_model_fn, config=config, train_batch_size=8, eval_batch_size=8,
    embedding_config_spec=embedding_config_spec)

TPUEstimator.train を呼び出して、モデルのトレーニングを開始します。

In [ ]:

estimator.train(_input_fn, steps=1)

次に、TPUEstimator.evaluate を呼び出して、評価データを使用してモデルを評価します。

In [ ]:

estimator.evaluate(_eval_input_fn, steps=1)

TensorFlow 2: TPUStrategy を使用して TPU の埋め込みをトレーニングする

TensorFlow 2 で TPU ワーカーをトレーニングするには、tf.distribute.TPUStrategy を Keras API とともに使用して、モデルの定義とトレーニング/評価を行います。（Keras Model.fit とカスタムトレーニングループ（tf.function と tf.GradientTape を使用）を使用したトレーニングのその他の例については、TPU の使用ガイドを参照してください。）

リモートクラスタに接続して TPU ワーカーを初期化するには初期化作業を行う必要があります。まず TPUClusterResolver を作成してクラスタ情報を提供し、クラスタに接続します。（詳しくは、TPU の使用ガイドの TPU 初期化 セクションを参照してください）。

In [ ]:

cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='')
tf.config.experimental_connect_to_cluster(cluster_resolver)
tf.tpu.experimental.initialize_tpu_system(cluster_resolver)
print("All devices: ", tf.config.list_logical_devices('TPU'))

次に、データを準備します。これは TensorFlow 1 の例でのデータセット作成方法と似ていますが、データセット関数に params ディクショナリではなく tf.distribute.InputContext オブジェクトが渡される点が異なります。このオブジェクトを使用して、ローカルバッチサイズ（およびこのパイプラインが対象とするホスト、データを適切に分割できるようにする）を決定できます。

tfrs.layers.embedding.TPUEmbedding API を使用する場合、Dataset.batch でデータセットをバッチ処理するときに drop_reminder=True オプションを含めることが重要です。TPUEmbedding には固定のバッチサイズが必要なためです。
さらに、評価とトレーニングが同じデバイスセットで行われている場合は、同じバッチサイズを使用する必要があります。
最後に、tf.keras.utils.experimental.DatasetCreator を特殊な入力オプション experimental_fetch_to_device=False と共に tf.distribute.InputOptions で使用する必要があります（ストラテジー固有の構成を保持します）。以下に例を示します

In [ ]:

global_batch_size = 8

def _input_dataset(context: tf.distribute.InputContext):
  dataset = tf.data.Dataset.from_tensor_slices((
      {"dense_feature": features,
       "sparse_feature": tf.SparseTensor(
           embedding_features_indices,
           embedding_features_values, [1, 2])},
           labels))
  dataset = dataset.shuffle(10).repeat()
  dataset = dataset.batch(
      context.get_per_replica_batch_size(global_batch_size),
      drop_remainder=True)
  return dataset.prefetch(2)

def _eval_dataset(context: tf.distribute.InputContext):
  dataset = tf.data.Dataset.from_tensor_slices((
      {"dense_feature": eval_features,
       "sparse_feature": tf.SparseTensor(
           eval_embedding_features_indices,
           eval_embedding_features_values, [1, 2])},
           eval_labels))
  dataset = dataset.repeat()
  dataset = dataset.batch(
      context.get_per_replica_batch_size(global_batch_size),
      drop_remainder=True)
  return dataset.prefetch(2)

input_options = tf.distribute.InputOptions(
    experimental_fetch_to_device=False)

input_dataset = tf.keras.utils.experimental.DatasetCreator(
    _input_dataset, input_options=input_options)

eval_dataset = tf.keras.utils.experimental.DatasetCreator(
    _eval_dataset, input_options=input_options)

次に、データの準備ができたら、TPUStrategy を作成し、このストラテジーの範囲（Strategy.scope）でモデル、指標、およびオプティマイザを定義します。

Model.compile で steps_per_execution の数値を選択する必要があります。これは、各 tf.function 呼び出し中に実行するバッチの数を指定するためです。これはパフォーマンスに重要です。この引数は、TPUEstimator で使用される iterations_per_loop に似ています。

tf.tpu.experimental.embedding_column（および tf.tpu.experimental.shared_embedding_column）を介して TensorFlow 1 で指定された特徴量と表の構成は、TensorFlow 2 では 2 つの構成オブジェクト直接指定できます。

tf.tpu.experimental.embedding.FeatureConfig
tf.tpu.experimental.embedding.TableConfig

（詳細については、関連する API ドキュメントを参照してください。）

In [ ]:

strategy = tf.distribute.TPUStrategy(cluster_resolver)
with strategy.scope():
  if hasattr(tf.keras.optimizers, "legacy"):
    optimizer = tf.keras.optimizers.legacy.Adagrad(learning_rate=0.05)
  else:
    optimizer = tf.keras.optimizers.Adagrad(learning_rate=0.05)
  dense_input = tf.keras.Input(shape=(2,), dtype=tf.float32, batch_size=global_batch_size)
  sparse_input = tf.keras.Input(shape=(), dtype=tf.int32, batch_size=global_batch_size)
  embedded_input = tfrs.layers.embedding.TPUEmbedding(
      feature_config=tf.tpu.experimental.embedding.FeatureConfig(
          table=tf.tpu.experimental.embedding.TableConfig(
              vocabulary_size=10,
              dim=5,
              initializer=tf.initializers.TruncatedNormal(mean=0.0, stddev=1)),
          name="sparse_input"),
      optimizer=optimizer)(sparse_input)
  input = tf.keras.layers.Concatenate(axis=1)([dense_input, embedded_input])
  result = tf.keras.layers.Dense(1)(input)
  model = tf.keras.Model(inputs={"dense_feature": dense_input, "sparse_feature": sparse_input}, outputs=result)
  model.compile(optimizer, "mse", steps_per_execution=10)

以上でトレーニングデータセットを使用してモデルをトレーニングする準備が整いました。

In [ ]:

model.fit(input_dataset, epochs=5, steps_per_epoch=10)

最後に、評価データセットを使用してモデルを評価します。

In [ ]:

model.evaluate(eval_dataset, steps=1, return_dict=True)

次のステップ

TPU 固有の埋め込みの設定の詳細については、API ドキュメントを参照してください。

tfrs.layers.embedding.TPUEmbedding: 特に特徴量と表の構成、オプティマイザの設定、モデルの作成（Keras functional API を使用するか、サブクラス化 tf.keras.Model を介して）、トレーニング/評価、および tf.saved_model を使用したモデルの提供
tf.tpu.experimental.embedding.TableConfig
tf.tpu.experimental.embedding.FeatureConfig

TensorFlow 2 の TPUStrategy の詳細については、次のリソースを参照してください。

ガイド: TPU の使用（Keras Model.fit を使用したトレーニング / tf.distribute.TPUStrategy を使用したカスタムトレーニングループと、tf.function でパフォーマンスを改善するためのヒントを含む）
ガイド: TensorFlow による分散トレーニング
ガイド: TPUEstimator から TPUStrategy に移行する。

トレーニングのカスタマイズの詳細については、次を参照してください。

ガイド: Model.fit の処理をカスタマイズする
ガイド: トレーニングループの新規作成

TPU（Google の機械学習専用 ASIC）は、Google Colab、TPU Research Cloud、Cloud TPU から入手できます。

Copyright 2021 The TensorFlow Authors.

TPU embedding_columns から TPUEmbedding レイヤーに移行する

セットアップ

TensorFlow 1: TPUEstimator を使用して TPU の埋め込みをトレーニングする

TensorFlow 2: TPUStrategy を使用して TPU の埋め込みをトレーニングする

次のステップ

Product

Resources

Company