tpu_estimator.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 TPUEstimator a TPUStrategy

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 sus flujos de trabajo que se ejecutan en TPUs desde la API TPUEstimator de TensorFlow 1 a la API TPUStrategy de TensorFlow 2.

• En TensorFlow 1, la API tf.compat.v1.estimator.tpu.TPUEstimator le permite entrenar y evaluar un modelo, así como realizar inferencias y guardar su modelo (para ser distribuido) en TPUs (en la nube).

• En TensorFlow 2, para realizar un entrenamiento síncrono en TPUs y TPU Pods (una recolección de dispositivos TPU conectados por interfaces de red dedicadas de alta velocidad), es necesario usar una estrategia de distribución de TPUs: tf.distribute.TPUStrategy. La estrategia puede funcionar con las API de Keras [incluyendo para la construcción de modelos (tf.keras.Model), optimizadores (tf.keras.optimizers.Optimizer) y entrenamiento (Model.fit)] así como con un bucle de entrenamiento personalizado (con tf.function y tf.GradientTape).

Para ver ejemplos de TensorFlow 2 de principio a fin, consulte la guía Utilizar TPUs (concretamente, la sección Clasificación en TPUs) y el tutorial Cómo resolver tareas GLUE utilizando BERT en TPU. También le puede resultar útil la guía Entrenamiento distribuido, que cubre todas las estrategias de distribución de TensorFlow, incluyendo TPUStrategy.

Preparación

Empieza con imports y un conjunto de datos sencillo a modo de demostración:

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

features = [[1., 1.5]]
labels = [[0.3]]
eval_features = [[4., 4.5]]
eval_labels = [[0.8]]

TensorFlow 1: Impulsar un modelo en TPUs con TPUEstimator

Esta sección de la guía muestra cómo realizar el entrenamiento y la evaluación con tf.compat.v1.estimator.tpu.TPUEstimator en TensorFlow 1.

Para usar un TPUEstimator, defina primero algunas funciones: una función de entrada para los datos de entrenamiento, una función de entrada de evaluación para los datos de evaluación y una función de modelo que indique al TPUEstimator 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((features, 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((eval_features, eval_labels))
  dataset = dataset.repeat()
  return dataset.batch(params['batch_size'], drop_remainder=True)

def _model_fn(features, labels, mode, params):
  logits = tf1.layers.Dense(1)(features)
  loss = tf1.losses.mean_squared_error(labels=labels, predictions=logits)
  optimizer = tf1.train.AdagradOptimizer(0.05)
  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 facilite la información del cluster, 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 guardado de puntos de verificación. Después, especificará el tamaño del lote tanto para el entrenamiento como para 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)
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)

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 la evaluación:

estimator.evaluate(_eval_input_fn, steps=1)

TensorFlow 2: Impulsar un modelo sobre TPUs con Keras Model.fit y 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, una vez preparados sus datos, creará una TPUStrategy, definirá un modelo, métricas y un optimizador en el ámbito de esta estrategia.

Para lograr una velocidad de entrenamiento comparable con TPUStrategy, procure elegir un número para steps_per_execution en Model.compile porque éste 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 un TPUEstimator. Si utiliza bucles de entrenamiento personalizados, debe asegurarse de que se ejecutan varios pasos dentro de la función de entrenamiento que usa tf.function. Vaya a la sección Mejorar el rendimiento con múltiples pasos dentro de tf.function de la guía Utilizar TPUs para obtener más información.

tf.distribute.TPUStrategy puede admitir formas dinámicas acotadas, en cuyo caso puede inferirse el límite superior del cálculo de la forma dinámica. Pero las formas dinámicas pueden suponer cierta sobrecarga de rendimiento en comparación con las formas estáticas. Así que, en general, se recomienda que las formas de entrada sean estáticas si es posible, especialmente en el entrenamiento. Una op común que devuelve una forma dinámica es tf.data.Dataset.batch(batch_size), ya que el número de muestras que quedan en un flujo puede ser inferior al tamaño del lote. Por lo tanto, al entrenar sobre la TPU, debería usar tf.data.Dataset.batch(..., drop_remainder=True) para obtener el mejor rendimiento en el entrenamiento.

dataset = tf.data.Dataset.from_tensor_slices(
    (features, labels)).shuffle(10).repeat().batch(
        8, drop_remainder=True).prefetch(2)
eval_dataset = tf.data.Dataset.from_tensor_slices(
    (eval_features, eval_labels)).batch(1, drop_remainder=True)

strategy = tf.distribute.TPUStrategy(cluster_resolver)
with strategy.scope():
  model = tf.keras.models.Sequential([tf.keras.layers.Dense(1)])
  optimizer = tf.keras.optimizers.Adagrad(learning_rate=0.05)
  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(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, return_dict=True)

Siguientes pasos

Para obtener más información sobre TPUStrategy en TensorFlow 2, tenga en cuenta 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

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.