 Voir sur TensorFlow.org |  Exécuter dans Google Colab |  Voir la source sur GitHub |  Télécharger le cahier |
Ce guide explique comment migrer vos workflows exécutés sur des TPU de l'API TPUStrategy de TensorFlow 1 vers l'API TPUEstimator de TensorFlow 2.
- Dans TensorFlow 1, l'API
tf.compat.v1.estimator.tpu.TPUEstimatorvous permet d'entraîner et d'évaluer un modèle, ainsi que d'effectuer des inférences et d'enregistrer votre modèle (pour diffusion) sur des TPU (Cloud). - Dans TensorFlow 2, pour effectuer un entraînement synchrone sur les TPU et les pods TPU (un ensemble d'appareils TPU connectés par des interfaces réseau haut débit dédiées), vous devez utiliser une stratégie de distribution TPU :
tf.distribute.TPUStrategy. La stratégie peut fonctionner avec les API Keras, y compris pour la création de modèles (tf.keras.Model), les optimiseurs (tf.keras.optimizers.Optimizer) et la formation (Model.fit) ainsi qu'une boucle de formation personnalisée (avectf.functionettf.GradientTape).
Pour des exemples TensorFlow 2 de bout en bout, consultez le guide Utiliser les TPU , à savoir la section Classification sur les TPU , et le didacticiel Résoudre les tâches GLUE à l'aide de BERT sur les TPU . Vous pouvez également trouver utile le guide de formation distribuée , qui couvre toutes les stratégies de distribution TensorFlow, y compris TPUStrategy .
Installer
Commencez par des importations et un jeu de données simple à des fins de démonstration :
import tensorflow as tf
import tensorflow.compat.v1 as tf1
/tmpfs/src/tf_docs_env/lib/python3.6/site-packages/requests/__init__.py:104: RequestsDependencyWarning: urllib3 (1.26.8) or chardet (2.3.0)/charset_normalizer (2.0.11) doesn't match a supported version! RequestsDependencyWarning)
features = [[1., 1.5]]
labels = [[0.3]]
eval_features = [[4., 4.5]]
eval_labels = [[0.8]]
TensorFlow 1 : Pilotez un modèle sur des TPU avec TPUEstimator
Cette section du guide explique comment effectuer une formation et une évaluation avec tf.compat.v1.estimator.tpu.TPUEstimator dans TensorFlow 1.
Pour utiliser un TPUEstimator , définissez d'abord quelques fonctions : une fonction d'entrée pour les données d'entraînement, une fonction d'entrée d'évaluation pour les données d'évaluation et une fonction de modèle qui indique à TPUEstimator comment l'opération d'entraînement est définie avec les fonctionnalités et les libellés :
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)
Une fois ces fonctions définies, créez un tf.distribute.cluster_resolver.TPUClusterResolver qui fournit les informations de cluster et un objet tf.compat.v1.estimator.tpu.RunConfig . En plus de la fonction de modèle que vous avez définie, vous pouvez maintenant créer un TPUEstimator . Ici, vous simplifierez le flux en sautant les économies de points de contrôle. Ensuite, vous spécifierez la taille du lot pour la formation et l'évaluation pour le TPUEstimator .
cluster_resolver = tf1.distribute.cluster_resolver.TPUClusterResolver(tpu='')
print("All devices: ", tf1.config.list_logical_devices('TPU'))
All devices: []
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)
WARNING:tensorflow:Estimator's model_fn (<function _model_fn at 0x7fef73ae76a8>) includes params argument, but params are not passed to Estimator.
WARNING:tensorflow:Using temporary folder as model directory: /tmp/tmp_bkua7zf
INFO:tensorflow:Using config: {'_model_dir': '/tmp/tmp_bkua7zf', '_tf_random_seed': None, '_save_summary_steps': 100, '_save_checkpoints_steps': None, '_save_checkpoints_secs': None, '_session_config': allow_soft_placement: true
cluster_def {
job {
name: "worker"
tasks {
key: 0
value: "10.240.1.2:8470"
}
}
}
isolate_session_state: true
, '_keep_checkpoint_max': 5, '_keep_checkpoint_every_n_hours': 10000, '_log_step_count_steps': None, '_train_distribute': None, '_device_fn': None, '_protocol': None, '_eval_distribute': None, '_experimental_distribute': None, '_experimental_max_worker_delay_secs': None, '_session_creation_timeout_secs': 7200, '_checkpoint_save_graph_def': True, '_service': None, '_cluster_spec': ClusterSpec({'worker': ['10.240.1.2:8470']}), '_task_type': 'worker', '_task_id': 0, '_global_id_in_cluster': 0, '_master': 'grpc://10.240.1.2:8470', '_evaluation_master': 'grpc://10.240.1.2:8470', '_is_chief': True, '_num_ps_replicas': 0, '_num_worker_replicas': 1, '_tpu_config': TPUConfig(iterations_per_loop=10, num_shards=None, num_cores_per_replica=None, per_host_input_for_training=2, tpu_job_name=None, initial_infeed_sleep_secs=None, input_partition_dims=None, eval_training_input_configuration=2, experimental_host_call_every_n_steps=1, experimental_allow_per_host_v2_parallel_get_next=False, experimental_feed_hook=None), '_cluster': <tensorflow.python.distribute.cluster_resolver.tpu.tpu_cluster_resolver.TPUClusterResolver object at 0x7ff288b6aa20>}
INFO:tensorflow:_TPUContext: eval_on_tpu True
Appelez TPUEstimator.train pour commencer à entraîner le modèle :
estimator.train(_input_fn, steps=1)
INFO:tensorflow:Querying Tensorflow master (grpc://10.240.1.2:8470) for TPU system metadata. INFO:tensorflow:Found TPU system: INFO:tensorflow:*** Num TPU Cores: 8 INFO:tensorflow:*** Num TPU Workers: 1 INFO:tensorflow:*** Num TPU Cores Per Worker: 8 INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:CPU:0, CPU, -1, 2562214468325910549) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:0, TPU, 17179869184, 7806191887455116208) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:1, TPU, 17179869184, 4935096526614797404) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:2, TPU, 17179869184, 6208852770722846295) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:3, TPU, 17179869184, -4484747666522931072) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:4, TPU, 17179869184, -8715412538518264422) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:5, TPU, 17179869184, -3521027846460785533) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:6, TPU, 17179869184, -6534172152637582552) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:7, TPU, 17179869184, 4735861352635655596) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU_SYSTEM:0, TPU_SYSTEM, 17179869184, -411508280321075475) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:XLA_CPU:0, XLA_CPU, 17179869184, 2431932884271560631) WARNING:tensorflow:From /tmpfs/src/tf_docs_env/lib/python3.6/site-packages/tensorflow/python/training/training_util.py:236: Variable.initialized_value (from tensorflow.python.ops.variables) is deprecated and will be removed in a future version. Instructions for updating: Use Variable.read_value. Variables in 2.X are initialized automatically both in eager and graph (inside tf.defun) contexts. INFO:tensorflow:Calling model_fn. WARNING:tensorflow:From /tmpfs/src/tf_docs_env/lib/python3.6/site-packages/tensorflow/python/training/adagrad.py:77: calling Constant.__init__ (from tensorflow.python.ops.init_ops) with dtype is deprecated and will be removed in a future version. Instructions for updating: Call initializer instance with the dtype argument instead of passing it to the constructor INFO:tensorflow:Bypassing TPUEstimator hook INFO:tensorflow:Done calling model_fn. INFO:tensorflow:TPU job name worker INFO:tensorflow:Graph was finalized. INFO:tensorflow:Running local_init_op. INFO:tensorflow:Done running local_init_op. WARNING:tensorflow:From /tmpfs/src/tf_docs_env/lib/python3.6/site-packages/tensorflow_estimator/python/estimator/tpu/tpu_estimator.py:758: Variable.load (from tensorflow.python.ops.variables) is deprecated and will be removed in a future version. Instructions for updating: Prefer Variable.assign which has equivalent behavior in 2.X. INFO:tensorflow:Initialized dataset iterators in 0 seconds INFO:tensorflow:Installing graceful shutdown hook. INFO:tensorflow:Creating heartbeat manager for ['/job:worker/replica:0/task:0/device:CPU:0'] INFO:tensorflow:Configuring worker heartbeat: shutdown_mode: WAIT_FOR_COORDINATOR INFO:tensorflow:Init TPU system INFO:tensorflow:Initialized TPU in 7 seconds INFO:tensorflow:Starting infeed thread controller. INFO:tensorflow:Starting outfeed thread controller. INFO:tensorflow:Enqueue next (1) batch(es) of data to infeed. INFO:tensorflow:Dequeue next (1) batch(es) of data from outfeed. INFO:tensorflow:Outfeed finished for iteration (0, 0) INFO:tensorflow:loss = 4.462118, step = 1 INFO:tensorflow:Stop infeed thread controller INFO:tensorflow:Shutting down InfeedController thread. INFO:tensorflow:InfeedController received shutdown signal, stopping. INFO:tensorflow:Infeed thread finished, shutting down. INFO:tensorflow:infeed marked as finished INFO:tensorflow:Stop output thread controller INFO:tensorflow:Shutting down OutfeedController thread. INFO:tensorflow:OutfeedController received shutdown signal, stopping. INFO:tensorflow:Outfeed thread finished, shutting down. INFO:tensorflow:outfeed marked as finished INFO:tensorflow:Shutdown TPU system. INFO:tensorflow:Loss for final step: 4.462118. INFO:tensorflow:training_loop marked as finished <tensorflow_estimator.python.estimator.tpu.tpu_estimator.TPUEstimator at 0x7fec59ef9d68>
Ensuite, appelez TPUEstimator.evaluate pour évaluer le modèle à l'aide des données d'évaluation :
estimator.evaluate(_eval_input_fn, steps=1)
INFO:tensorflow:Could not find trained model in model_dir: /tmp/tmp_bkua7zf, running initialization to evaluate.
INFO:tensorflow:Calling model_fn.
WARNING:tensorflow:From /tmpfs/src/tf_docs_env/lib/python3.6/site-packages/tensorflow_estimator/python/estimator/tpu/tpu_estimator.py:3406: div (from tensorflow.python.ops.math_ops) is deprecated and will be removed in a future version.
Instructions for updating:
Deprecated in favor of operator or tf.math.divide.
INFO:tensorflow:Done calling model_fn.
INFO:tensorflow:Starting evaluation at 2022-02-05T13:15:25
INFO:tensorflow:TPU job name worker
INFO:tensorflow:Graph was finalized.
INFO:tensorflow:Running local_init_op.
INFO:tensorflow:Done running local_init_op.
INFO:tensorflow:Init TPU system
INFO:tensorflow:Initialized TPU in 10 seconds
INFO:tensorflow:Starting infeed thread controller.
INFO:tensorflow:Starting outfeed thread controller.
INFO:tensorflow:Initialized dataset iterators in 0 seconds
INFO:tensorflow:Enqueue next (1) batch(es) of data to infeed.
INFO:tensorflow:Dequeue next (1) batch(es) of data from outfeed.
INFO:tensorflow:Outfeed finished for iteration (0, 0)
INFO:tensorflow:Evaluation [1/1]
INFO:tensorflow:Stop infeed thread controller
INFO:tensorflow:Shutting down InfeedController thread.
INFO:tensorflow:InfeedController received shutdown signal, stopping.
INFO:tensorflow:Infeed thread finished, shutting down.
INFO:tensorflow:infeed marked as finished
INFO:tensorflow:Stop output thread controller
INFO:tensorflow:Shutting down OutfeedController thread.
INFO:tensorflow:OutfeedController received shutdown signal, stopping.
INFO:tensorflow:Outfeed thread finished, shutting down.
INFO:tensorflow:outfeed marked as finished
INFO:tensorflow:Shutdown TPU system.
INFO:tensorflow:Inference Time : 10.80091s
INFO:tensorflow:Finished evaluation at 2022-02-05-13:15:36
INFO:tensorflow:Saving dict for global step 1: global_step = 1, loss = 116.58184
INFO:tensorflow:evaluation_loop marked as finished
{'loss': 116.58184, 'global_step': 1}
TensorFlow 2 : Pilotez un modèle sur des TPU avec Keras Model.fit et TPUStrategy
Dans TensorFlow 2, pour former les nœuds de calcul TPU, utilisez tf.distribute.TPUStrategy avec les API Keras pour la définition du modèle et la formation/l'évaluation. (Reportez-vous au guide Utiliser les TPU pour plus d'exemples d'entraînement avec Keras Model.fit et une boucle d'entraînement personnalisée (avec tf.function et tf.GradientTape ).)
Étant donné que vous devez effectuer un travail d'initialisation pour vous connecter au cluster distant et initialiser les travailleurs TPU, commencez par créer un TPUClusterResolver pour fournir les informations de cluster et vous connecter au cluster. (Pour en savoir plus, consultez la section d' initialisation des TPU du guide Utiliser les TPU.)
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'))
INFO:tensorflow:Clearing out eager caches INFO:tensorflow:Clearing out eager caches INFO:tensorflow:Initializing the TPU system: grpc://10.240.1.2:8470 INFO:tensorflow:Initializing the TPU system: grpc://10.240.1.2:8470 INFO:tensorflow:Finished initializing TPU system. INFO:tensorflow:Finished initializing TPU system. All devices: [LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:0', device_type='TPU'), LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:1', device_type='TPU'), LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:2', device_type='TPU'), LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:3', device_type='TPU'), LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:4', device_type='TPU'), LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:5', device_type='TPU'), LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:6', device_type='TPU'), LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:7', device_type='TPU')]
Ensuite, une fois vos données préparées, vous allez créer une TPUStrategy , définir un modèle, des métriques et un optimiseur dans le cadre de cette stratégie.
Pour obtenir une vitesse d'entraînement comparable avec TPUStrategy , vous devez vous assurer de choisir un nombre pour steps_per_execution dans Model.compile car il spécifie le nombre de lots à exécuter lors de chaque appel tf.function et est essentiel pour les performances. Cet argument est similaire à iterations_per_loop utilisé dans un TPUEstimator . Si vous utilisez des boucles d'entraînement personnalisées, vous devez vous assurer que plusieurs étapes sont exécutées dans la fonction d'entraînement tf.function -ed. Accédez à la section Améliorer les performances en plusieurs étapes dans tf.function du guide Utiliser les TPU pour plus d'informations.
tf.distribute.TPUStrategy peut prendre en charge des formes dynamiques limitées, c'est-à-dire que la limite supérieure du calcul de la forme dynamique peut être déduite. Mais les formes dynamiques peuvent introduire une surcharge de performances par rapport aux formes statiques. Ainsi, il est généralement recommandé de rendre vos formes d'entrée statiques si possible, en particulier lors de la formation. Une opération courante qui renvoie une forme dynamique est tf.data.Dataset.batch(batch_size) , car le nombre d'échantillons restant dans un flux peut être inférieur à la taille du lot. Par conséquent, lors de l'entraînement sur le TPU, vous devez utiliser tf.data.Dataset.batch(..., drop_remainder=True) pour de meilleures performances d'entraînement.
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)
INFO:tensorflow:Found TPU system: INFO:tensorflow:Found TPU system: INFO:tensorflow:*** Num TPU Cores: 8 INFO:tensorflow:*** Num TPU Cores: 8 INFO:tensorflow:*** Num TPU Workers: 1 INFO:tensorflow:*** Num TPU Workers: 1 INFO:tensorflow:*** Num TPU Cores Per Worker: 8 INFO:tensorflow:*** Num TPU Cores Per Worker: 8 INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:localhost/replica:0/task:0/device:CPU:0, CPU, 0, 0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:localhost/replica:0/task:0/device:CPU:0, CPU, 0, 0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:CPU:0, CPU, 0, 0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:CPU:0, CPU, 0, 0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:0, TPU, 0, 0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:0, TPU, 0, 0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:1, TPU, 0, 0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:1, TPU, 0, 0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:2, TPU, 0, 0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:2, TPU, 0, 0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:3, TPU, 0, 0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:3, TPU, 0, 0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:4, TPU, 0, 0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:4, TPU, 0, 0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:5, TPU, 0, 0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:5, TPU, 0, 0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:6, TPU, 0, 0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:6, TPU, 0, 0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:7, TPU, 0, 0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:7, TPU, 0, 0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU_SYSTEM:0, TPU_SYSTEM, 0, 0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU_SYSTEM:0, TPU_SYSTEM, 0, 0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:XLA_CPU:0, XLA_CPU, 0, 0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:XLA_CPU:0, XLA_CPU, 0, 0)
Avec cela, vous êtes prêt à entraîner le modèle avec l'ensemble de données d'entraînement :
model.fit(dataset, epochs=5, steps_per_epoch=10)
Epoch 1/5 10/10 [==============================] - 2s 151ms/step - loss: 0.0840 Epoch 2/5 10/10 [==============================] - 0s 3ms/step - loss: 9.6915e-04 Epoch 3/5 10/10 [==============================] - 0s 3ms/step - loss: 1.5100e-05 Epoch 4/5 10/10 [==============================] - 0s 3ms/step - loss: 2.3593e-07 Epoch 5/5 10/10 [==============================] - 0s 3ms/step - loss: 3.7059e-09 <keras.callbacks.History at 0x7fec58275438>
Enfin, évaluez le modèle à l'aide de l'ensemble de données d'évaluation :
model.evaluate(eval_dataset, return_dict=True)
1/1 [==============================] - 2s 2s/step - loss: 0.6127
{'loss': 0.6127181053161621}
Prochaines étapes
Pour en savoir plus sur TPUStrategy dans TensorFlow 2, consultez les ressources suivantes :
- Guide : Utiliser des TPU (couvrant l'entraînement avec Keras
Model.fit/une boucle d'entraînement personnalisée avectf.distribute.TPUStrategy, ainsi que des conseils pour améliorer les performances avectf.function) - Guide : Entraînement distribué avec TensorFlow
Pour en savoir plus sur la personnalisation de votre formation, consultez :
- Guide : Personnalisez ce qui se passe dans Model.fit
- Guide : Rédaction d'une boucle d'entraînement à partir de zéro
Les TPU, les ASIC spécialisés de Google pour le machine learning, sont disponibles via Google Colab , le TPU Research Cloud et Cloud TPU .