mirrored_strategy.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 el entrenamiento monotrabajador con múltiples GPUs

Ver en TensorFlow.org

Ejecutar en Google Colab

Ver código fuente en GitHub

Descargar bloc de notas

Esta guía muestra cómo migrar los flujos de trabajo monotrabajador y con múltiples GPU de TensorFlow 1 a TensorFlow 2.

Para realizar un entrenamiento síncrono en múltiples GPUs de una misma máquina:

• En TensorFlow 1, usas las API tf.estimator.Estimator con tf.distribute.MirroredStrategy.

• En TensorFlow 2, puedes usar Model.fit de Keras o un bucle de entrenamiento personalizado con tf.distribute.MirroredStrategy. Aprende más en la guía Entrenamiento distribuido con TensorFlow.

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], [2., 2.5], [3., 3.5]]
labels = [[0.3], [0.5], [0.7]]
eval_features = [[4., 4.5], [5., 5.5], [6., 6.5]]
eval_labels = [[0.8], [0.9], [1.]]

TensorFlow 1: Entrenamiento distribuido monotrabajador con tf.estimator.Estimator

Este ejemplo demuestra el flujo de trabajo canónico de TensorFlow 1 de entrenamiento monotrabajador con múltiples GPU. Tienes que configurar la estrategia de distribución (tf.distribute.MirroredStrategy) a través del parámetro config del tf.estimator.Estimator:

def _input_fn():
  return tf1.data.Dataset.from_tensor_slices((features, labels)).batch(1)

def _eval_input_fn():
  return tf1.data.Dataset.from_tensor_slices(
      (eval_features, eval_labels)).batch(1)

def _model_fn(features, labels, mode):
  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.EstimatorSpec(mode, loss=loss, train_op=train_op)

strategy = tf1.distribute.MirroredStrategy()
config = tf1.estimator.RunConfig(
    train_distribute=strategy, eval_distribute=strategy)
estimator = tf1.estimator.Estimator(model_fn=_model_fn, config=config)

train_spec = tf1.estimator.TrainSpec(input_fn=_input_fn)
eval_spec = tf1.estimator.EvalSpec(input_fn=_eval_input_fn)
tf1.estimator.train_and_evaluate(estimator, train_spec, eval_spec)

TensorFlow 2: Entrenamiento monotrabajador con Keras

Al migrar a TensorFlow 2, puedes usar las API de Keras con tf.distribute.MirroredStrategy.

Si usas las API tf.keras para construir el modelo y Model.fit de Keras para el entrenamiento, la principal diferencia es instanciar el modelo Keras, un optimizador y métricas en el contexto de Strategy.scope, en lugar de definir un config para tf.estimator.Estimator.

Si necesitas usar un bucle de entrenamiento personalizado, consulta la guía Usar tf.distribute.Strategy con bucles de entrenamiento personalizados.

dataset = tf.data.Dataset.from_tensor_slices((features, labels)).batch(1)
eval_dataset = tf.data.Dataset.from_tensor_slices(
      (eval_features, eval_labels)).batch(1)

strategy = tf.distribute.MirroredStrategy()
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=optimizer, loss='mse')
model.fit(dataset)
model.evaluate(eval_dataset, return_dict=True)

Siguientes pasos

Para obtener más información sobre el entrenamiento distribuido con tf.distribute.MirroredStrategy en TensorFlow 2, consulta la siguiente documentación:

• El tutorial Entrenamiento distribuido en una máquina con Keras

• El tutorial Entrenamiento distribuido en una máquina con un bucle de entrenamiento personalizado

• La guía Entrenamiento distribuido con TensorFlow

• La guía Cómo usar múltiples GPU

• La guía Optimizar el rendimiento en el host único multi-GPU (con el Perfilador TensorFlow)