tpu_embedding.ipynb

Kernel: Python 3

Copyright 2021 The TensorFlow Authors.

#@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.

Migrar de TPU embedding_columns a capa TPUEmbedding

Ver en TensorFlow.org

Ejecutar en Google Colab

Ver código fuente en GitHub

Descargar bloc de notas

Esta guía demuestra cómo migrar el entrenamiento de incorporación en TPUs de la API embedding_column de TensorFlow 1 con TPUEstimator a la API de la capa TPUEmbedding de TensorFlow 2 con TPUStrategy.

Las incorporaciones son matrices (grandes). Son tablas de búsqueda que mapean a partir de un espacio de características disperso a vectores densas. Las incorporaciones ofrecen representaciones eficientes y densas que captan las similitudes y relaciones complejas entre las características.

TensorFlow incluye soporte especializado para el entrenamiento de incorporaciones en TPUs. Este soporte de incorporación específico de TPU le permite entrenar incorporaciones que son más grandes que la memoria de un solo dispositivo TPU, y usar entradas dispersas y desiguales en TPUs.

• En TensorFlow 1, tf.compat.v1.estimator.tpu.TPUEstimator es una API de alto nivel que encapsula el entrenamiento, la evaluación, la predicción y la exportación para su servicio con TPUs. Tiene un soporte especial para tf.compat.v1.tpu.experimental.embedding_column.

• Para implementarlo en TensorFlow 2, use la capa tfrs.layers.embedding.TPUEmbedding de TensorFlow Recommenders. Para el entrenamiento y la evaluación, use una estrategia de distribución TPU (tf.distribute.TPUStrategy) que sea compatible con las API de Keras para, por ejemplo, la construcción de modelos (tf.keras. Model), optimizadores (tf.keras.optimizers.Optimizer), y entrenamiento con Model.fit o un bucle de entrenamiento personalizado con tf.function y tf.GradientTape.

Para más información, consulte la documentación de la API de la capa tfrs.layers.embedding.TPUEmbedding, así como los documentos de tf.tpu.experimental.embedding.TableConfig y tf.tpu.experimental.embedding.FeatureConfig. Para una visión general de tf.distribute.TPUStrategy, consulte la guía Entrenamiento distribuido y la guía Utilizar TPUs. Si está migrando de TPUEstimator a TPUStrategy, consulte la guía de migración de TPU.

Preparación

Comience instalando TensorFlow Recommenders e importando algunos paquetes necesarios:

!pip install tensorflow-recommenders

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

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

Y prepare un conjunto de datos sencillo para fines de demostración:

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: Entrene incorporaciones en TPUs con TPUEstimator

En TensorFlow 1, se preparan incorporaciones TPU usando la API tf.compat.v1.tpu.experimental.embedding_column y entrene/evalúe el modelo sobre TPUs con tf.compat.v1.estimator.tpu.TPUEstimator.

Las entradas son números enteros que van de cero al tamaño del vocabulario para la tabla de incorporación TPU. Comience por codificar las entradas en ID categórica con tf.feature_column.categorical_column_with_identity. Use "sparse_feature" para el parámetro key, ya que las características de entrada son de valor entero, mientras que num_buckets es el tamaño del vocabulario para la tabla de incorporación (10).

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

A continuación, convierta las entradas categóricas dispersas en una representación densa con tpu.experimental.embedding_column, donde dimension es la anchura de la tabla de incorporación. Almacenará un vector de incorporación para cada uno de los num_buckets.

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

Ahora, defina la configuración de incorporación específica de la TPU mediante tf.estimator.tpu.experimental.EmbeddingConfigSpec. Más tarde se la pasará a tf.estimator.tpu.TPUEstimator como parámetro embedding_config_spec.

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

A continuación, para usar un TPUEstimator, defina:

• Una función de entrada para los datos de entrenamiento

• Una función de entrada de evaluación para los datos de evaluación

• Una función modelo para dar instrucciones al TPUEstimator sobre cómo se define la op de entrenamiento con las características y las etiquetas.

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)

Con esas funciones definidas, cree un objeto tf.distribute.cluster_resolver.TPUClusterResolver que ofrezca la información del clúster, y un objeto tf.compat.v1.estimator.tpu.RunConfig.

Junto con la función modelo que ha definido, ahora puede crear un TPUEstimator. Aquí, simplificará el flujo omitiendo el ahorro de puntos de verificación. A continuación, especificará el tamaño del lote tanto para el entrenamiento como en la evaluación para el TPUEstimator.

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

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)

Llame a TPUEstimator.train para comenzar el entrenamiento del modelo:

estimator.train(_input_fn, steps=1)

Luego, llame a TPUEstimator.evaluate para evaluar el modelo usando los datos de evaluación:

estimator.evaluate(_eval_input_fn, steps=1)

TensorFlow 2: Entrene incorporaciones en TPUs con TPUStrategy

En TensorFlow 2, para entrenar sobre los trabajadores TPU, use tf.distribute.TPUStrategy junto con las APIs de Keras para la definición y entrenamiento / evaluación del modelo (consulte la guía Utilizar TPU para ver más ejemplos de entrenamiento con Model.fit de Keras y un bucle de entrenamiento personalizado (con tf.function y tf.GradientTape)).

Dado que necesita realizar algún trabajo de inicialización para conectarse al cluster remoto e inicializar los trabajadores de la TPU, comience creando un TPUClusterResolver para proveer la información del cluster y conectarse al mismo. Puede aprender más en la sección Inicialización de TPU de la guía Utilizar TPUs.

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'))

A continuación, prepare sus datos. Esta operación es similar a cómo creó un conjunto de datos en el ejemplo de TensorFlow 1, excepto que ahora se pasa a la función del conjunto de datos un objeto tf.distribute.InputContext en lugar de un dict params. Puede usar este objeto para determinar el tamaño del lote local (y para qué host es esta canalización, de modo que pueda particionar adecuadamente sus datos).

• Al usar la API tfrs.layers.embedding.TPUEmbedding, es importante incluir la opción drop_remainder=True al procesar por lotes el conjunto de datos con Dataset.batch, ya que TPUEmbedding requiere un tamaño de lote fijo.

• Además, debe usarse el mismo tamaño de lote para la evaluación y el entrenamiento si tienen lugar en el mismo conjunto de dispositivos.

• Por último, debe usar tf.keras.utils.experimental.DatasetCreator junto con la opción de entrada especial (experimental_fetch_to_device=False) en tf.distribute.InputOptions (que alberga configuraciones específicas de la estrategia). Esto se demuestra a continuación:

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)

Después, una vez preparados los datos, creará una TPUStrategy, y definirá un modelo, métricas y un optimizador bajo el ámbito de esta estrategia (Strategy.scope).

Debería elegir un número para steps_per_execution en Model.compile ya que especifica el número de lotes a ejecutar durante cada llamada a tf.function, y es crítico para el rendimiento. Este argumento es similar a iterations_per_loop usado en TPUEstimator.

Las características y la configuración de la tabla que se especificaron en TensorFlow 1 a través de tf.tpu.experimental.embedding_column (y tf.tpu.experimental.shared_embedding_column) pueden especificarse directamente en TensorFlow 2 a través de un par de objetos de configuración:

• tf.tpu.experimental.embedding.FeatureConfig

• tf.tpu.experimental.embedding.TableConfig

(Consulte la documentación de la API asociada para ver más detalles).

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)

Una vez hecho esto, ya está listo para entrenar el modelo con el conjunto de datos de entrenamiento:

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

Por último, evalúe el modelo usando el conjunto de datos de evaluación:

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

Siguientes pasos

Encontrará más información sobre cómo configurar las incorporaciones específicas de TPU en los documentos de la API:

• tfrs.layers.embedding.TPUEmbedding : particularmente sobre la configuración de funciones y tablas, configuración del optimizador, creación de un modelo (usando la API funcional de Keras o mediante subclases tf.keras.Model), entrenamiento/evaluación y entrega de modelos con tf.saved_model

• tf.tpu.experimental.embedding.TableConfig

• tf.tpu.experimental.embedding.FeatureConfig

Para más información sobre TPUStrategy en TensorFlow 2, considere los siguientes recursos:

• Guía: Utilizar TPUs (que cubre el entrenamiento con Keras Model.fit/un bucle de entrenamiento personalizado con tf.distribute.TPUStrategy, así como consejos para mejorar el rendimiento con tf.function)

• Guía: Entrenamiento distribuido con TensorFlow

• Guía: Migrar de TPUEstimator a TPUStrategy

Para saber más sobre cómo personalizar su entrenamiento, consulte:

• Guía: Personalice lo que ocurre en Model.fit

• Guía: Escribir un bucle de entrenamiento desde cero

Las TPU (ASIC especializadas de Google para el aprendizaje automático) están disponibles a través de Google Colab, la Cloud TPU Research Cloud y Cloud TPU.