Formation distribuée avec TensorFlow

Voir sur TensorFlow.org Exécuter dans Google Colab Voir la source sur GitHub Télécharger le cahier

Aperçu

tf.distribute.Strategy est une API TensorFlow permettant de répartir la formation sur plusieurs GPU, plusieurs machines ou TPU. À l'aide de cette API, vous pouvez distribuer vos modèles existants et votre code de formation avec des modifications de code minimales.

tf.distribute.Strategy a été conçu avec ces objectifs clés à l'esprit :

  • Facile à utiliser et prenant en charge plusieurs segments d'utilisateurs, y compris les chercheurs, les ingénieurs en apprentissage automatique, etc.
  • Fournit de bonnes performances prêtes à l'emploi.
  • Commutation facile entre les stratégies.

Vous pouvez distribuer la formation à l'aide de tf.distribute.Strategy avec une API de haut niveau telle que Keras Model.fit , ainsi que des boucles de formation personnalisées (et, en général, tout calcul utilisant TensorFlow).

Dans TensorFlow 2.x, vous pouvez exécuter vos programmes avec empressement ou dans un graphique à l'aide tf.function . tf.distribute.Strategy a l'intention de prendre en charge ces deux modes d'exécution, mais fonctionne mieux avec tf.function . Le mode Eager est uniquement recommandé à des fins de débogage et n'est pas pris en charge pour tf.distribute.TPUStrategy . Bien que la formation soit au centre de ce guide, cette API peut également être utilisée pour distribuer l'évaluation et la prédiction sur différentes plates-formes.

Vous pouvez utiliser tf.distribute.Strategy avec très peu de modifications de votre code, car les composants sous-jacents de TensorFlow ont été modifiés pour devenir sensibles à la stratégie. Cela inclut les variables, les couches, les modèles, les optimiseurs, les métriques, les résumés et les points de contrôle.

Dans ce guide, vous découvrirez différents types de stratégies et comment vous pouvez les utiliser dans différentes situations. Pour savoir comment déboguer les problèmes de performances, consultez le guide des performances GPU Optimize TensorFlow .

Configurer TensorFlow

import tensorflow as tf

Types de stratégies

tf.distribute.Strategy entend couvrir un certain nombre de cas d'utilisation selon différents axes. Certaines de ces combinaisons sont actuellement prises en charge et d'autres seront ajoutées à l'avenir. Certains de ces axes sont :

  • Entraînement synchrone ou asynchrone : il s'agit de deux manières courantes de distribuer l'entraînement avec le parallélisme des données. Dans la formation synchronisée, tous les travailleurs s'entraînent sur différentes tranches de données d'entrée synchronisées et agrégent des gradients à chaque étape. Dans la formation asynchrone, tous les travailleurs s'entraînent indépendamment sur les données d'entrée et mettent à jour les variables de manière asynchrone. En règle générale, la formation de synchronisation est prise en charge via all-reduce et async via l'architecture du serveur de paramètres.
  • Plate-forme matérielle : vous souhaiterez peut-être adapter votre formation à plusieurs GPU sur une seule machine, ou à plusieurs machines dans un réseau (avec 0 GPU ou plus chacune), ou sur des Cloud TPU.

Afin de prendre en charge ces cas d'utilisation, TensorFlow a MirroredStrategy , TPUStrategy , MultiWorkerMirroredStrategy , ParameterServerStrategy , CentralStorageStrategy , ainsi que d'autres stratégies disponibles. La section suivante explique lesquels d'entre eux sont pris en charge dans quels scénarios dans TensorFlow. Voici un aperçu rapide :

API de formation MirroredStrategy TPUStrategy MultiWorkerMirroredStrategy CentralStorageStrategy ParameterServerStrategy
Keras Model.fit Prise en charge Prise en charge Prise en charge Support expérimental Support expérimental
Boucle d'entraînement personnalisée Prise en charge Prise en charge Prise en charge Support expérimental Support expérimental
API d'estimation Assistance limitée Non supporté Assistance limitée Assistance limitée Assistance limitée

Stratégie en miroir

tf.distribute.MirroredStrategy prend en charge la formation distribuée synchrone sur plusieurs GPU sur une seule machine. Il crée une réplique par périphérique GPU. Chaque variable du modèle est reflétée dans toutes les répliques. Ensemble, ces variables forment une seule variable conceptuelle appelée MirroredVariable . Ces variables sont maintenues synchronisées les unes avec les autres en appliquant des mises à jour identiques.

Des algorithmes all-reduce efficaces sont utilisés pour communiquer les mises à jour des variables entre les appareils. All-reduce regroupe les tenseurs sur tous les appareils en les additionnant et les rend disponibles sur chaque appareil. C'est un algorithme fusionné qui est très efficace et peut réduire considérablement la surcharge de synchronisation. Il existe de nombreux algorithmes et implémentations tout-réduits disponibles, selon le type de communication disponible entre les appareils. Par défaut, il utilise la NVIDIA Collective Communication Library ( NCCL ) comme implémentation all-reduce. Vous pouvez choisir parmi quelques autres options ou écrire la vôtre.

Voici la manière la plus simple de créer MirroredStrategy :

mirrored_strategy = tf.distribute.MirroredStrategy()
INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:GPU:0',)

Cela créera une instance MirroredStrategy , qui utilisera tous les GPU visibles pour TensorFlow et NCCL, en tant que communication entre appareils.

Si vous souhaitez n'utiliser que certains GPU sur votre machine, vous pouvez le faire comme ceci :

mirrored_strategy = tf.distribute.MirroredStrategy(devices=["/gpu:0", "/gpu:1"])
WARNING:tensorflow:Some requested devices in `tf.distribute.Strategy` are not visible to TensorFlow: /job:localhost/replica:0/task:0/device:GPU:1,/job:localhost/replica:0/task:0/device:GPU:0
INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:GPU:0', '/job:localhost/replica:0/task:0/device:GPU:1')

Si vous souhaitez remplacer la communication entre appareils, vous pouvez le faire à l'aide de l'argument cross_device_ops en fournissant une instance de tf.distribute.CrossDeviceOps . Actuellement, tf.distribute.HierarchicalCopyAllReduce et tf.distribute.ReductionToOneDevice sont deux options autres que tf.distribute.NcclAllReduce , qui est la valeur par défaut.

mirrored_strategy = tf.distribute.MirroredStrategy(
    cross_device_ops=tf.distribute.HierarchicalCopyAllReduce())
INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:GPU:0',)

TPUStratégie

tf.distribute.TPUStrategy vous permet d'exécuter votre formation TensorFlow sur des unités de traitement Tensor (TPU) . Les TPU sont des ASIC spécialisés de Google conçus pour accélérer considérablement les charges de travail d'apprentissage automatique. Ils sont disponibles sur Google Colab , le TPU Research Cloud et Cloud TPU .

En termes d'architecture de formation distribuée, TPUStrategy est identique à MirroredStrategy - il implémente une formation distribuée synchrone. Les TPU fournissent leur propre implémentation d'opérations efficaces de réduction globale et d'autres opérations collectives sur plusieurs cœurs de TPU, qui sont utilisées dans TPUStrategy .

Voici comment instancier TPUStrategy :

cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver(
    tpu=tpu_address)
tf.config.experimental_connect_to_cluster(cluster_resolver)
tf.tpu.experimental.initialize_tpu_system(cluster_resolver)
tpu_strategy = tf.distribute.TPUStrategy(cluster_resolver)

L'instance TPUClusterResolver permet de localiser les TPU. Dans Colab, vous n'avez pas besoin de spécifier d'arguments.

Si vous souhaitez l'utiliser pour les Cloud TPU :

  • Vous devez spécifier le nom de votre ressource TPU dans l'argument tpu .
  • Vous devez initialiser explicitement le système TPU au démarrage du programme. Ceci est nécessaire avant que les TPU puissent être utilisés pour le calcul. L'initialisation du système TPU efface également la mémoire du TPU, il est donc important de terminer cette étape en premier afin d'éviter de perdre l'état.

MultiWorkerMirroredStrategy

tf.distribute.MultiWorkerMirroredStrategy est très similaire à MirroredStrategy . Il implémente une formation distribuée synchrone sur plusieurs travailleurs, chacun avec potentiellement plusieurs GPU. Semblable à tf.distribute.MirroredStrategy , il crée des copies de toutes les variables du modèle sur chaque appareil sur tous les travailleurs.

Voici la manière la plus simple de créer MultiWorkerMirroredStrategy :

strategy = tf.distribute.MultiWorkerMirroredStrategy()
WARNING:tensorflow:Collective ops is not configured at program startup. Some performance features may not be enabled.
INFO:tensorflow:Single-worker MultiWorkerMirroredStrategy with local_devices = ('/device:GPU:0',), communication = CommunicationImplementation.AUTO

MultiWorkerMirroredStrategy a deux implémentations pour les communications entre appareils. CommunicationImplementation.RING est basé sur RPC et prend en charge les CPU et les GPU. CommunicationImplementation.NCCL utilise NCCL et fournit des performances de pointe sur les GPU, mais il ne prend pas en charge les CPU. CollectiveCommunication.AUTO reporte le choix à Tensorflow. Vous pouvez les spécifier de la manière suivante :

communication_options = tf.distribute.experimental.CommunicationOptions(
    implementation=tf.distribute.experimental.CommunicationImplementation.NCCL)
strategy = tf.distribute.MultiWorkerMirroredStrategy(
    communication_options=communication_options)
WARNING:tensorflow:Collective ops is not configured at program startup. Some performance features may not be enabled.
INFO:tensorflow:Single-worker MultiWorkerMirroredStrategy with local_devices = ('/device:GPU:0',), communication = CommunicationImplementation.NCCL

L'une des principales différences entre la formation multi-travailleurs et la formation multi-GPU est la configuration multi-travailleurs. La variable d'environnement 'TF_CONFIG' est la méthode standard dans TensorFlow pour spécifier la configuration du cluster à chaque nœud de calcul faisant partie du cluster. Pour en savoir plus , consultez la section Configuration de TF_CONFIG de ce document.

Pour plus de détails sur MultiWorkerMirroredStrategy , consultez les didacticiels suivants :

ParameterServerStrategy

La formation de serveur de paramètres est une méthode commune de données parallèles pour mettre à l'échelle la formation de modèles sur plusieurs machines. Un cluster de formation de serveur de paramètres se compose de nœuds de calcul et de serveurs de paramètres. Les variables sont créées sur les serveurs de paramètres et elles sont lues et mises à jour par les travailleurs à chaque étape. Consultez le didacticiel de formation sur le serveur de paramètres pour plus de détails.

Dans TensorFlow 2, la formation du serveur de paramètres utilise une architecture centrale basée sur un coordinateur via la classe tf.distribute.experimental.coordinator.ClusterCoordinator .

Dans cette implémentation, les tâches de worker et de parameter server exécutent tf.distribute.Server qui écoutent les tâches du coordinateur. Le coordinateur crée des ressources, répartit les tâches de formation, rédige des points de contrôle et traite les échecs de tâche.

Dans la programmation exécutée sur le coordinateur, vous utiliserez un objet ParameterServerStrategy pour définir une étape de formation et utiliserez un ClusterCoordinator pour distribuer les étapes de formation aux travailleurs distants. Voici la manière la plus simple de les créer :

strategy = tf.distribute.experimental.ParameterServerStrategy(
    tf.distribute.cluster_resolver.TFConfigClusterResolver(),
    variable_partitioner=variable_partitioner)
coordinator = tf.distribute.experimental.coordinator.ClusterCoordinator(
    strategy)

Pour en savoir plus sur ParameterServerStrategy , consultez la formation sur le serveur de paramètres avec Keras Model.fit et un didacticiel de boucle de formation personnalisé .

Dans TensorFlow 1, ParameterServerStrategy est disponible uniquement avec un estimateur via le symbole tf.compat.v1.distribute.experimental.ParameterServerStrategy .

Stratégie de stockage central

tf.distribute.experimental.CentralStorageStrategy effectue également une formation synchrone. Les variables ne sont pas mises en miroir, mais placées sur le CPU et les opérations sont répliquées sur tous les GPU locaux. S'il n'y a qu'un seul GPU, toutes les variables et opérations seront placées sur ce GPU.

Créez une instance de CentralStorageStrategy en :

central_storage_strategy = tf.distribute.experimental.CentralStorageStrategy()
INFO:tensorflow:ParameterServerStrategy (CentralStorageStrategy if you are using a single machine) with compute_devices = ['/job:localhost/replica:0/task:0/device:GPU:0'], variable_device = '/job:localhost/replica:0/task:0/device:GPU:0'

Cela créera une instance CentralStorageStrategy qui utilisera tous les GPU et CPU visibles. La mise à jour des variables sur les répliques sera agrégée avant d'être appliquée aux variables.

Autres stratégies

En plus des stratégies ci-dessus, il existe deux autres stratégies qui pourraient être utiles pour le prototypage et le débogage lors de l'utilisation des API tf.distribute .

Stratégie par défaut

La stratégie par défaut est une stratégie de distribution qui est présente lorsqu'aucune stratégie de distribution explicite n'est dans la portée. Il implémente l'interface tf.distribute.Strategy mais est un passe-système et ne fournit aucune distribution réelle. Par exemple, Strategy.run(fn) appellera simplement fn . Le code écrit à l'aide de cette stratégie doit se comporter exactement comme le code écrit sans aucune stratégie. Vous pouvez considérer cela comme une stratégie "no-op".

La stratégie par défaut est un singleton et on ne peut pas en créer d'autres instances. Il peut être obtenu en utilisant tf.distribute.get_strategy en dehors de la portée de toute stratégie explicite (la même API qui peut être utilisée pour obtenir la stratégie actuelle à l'intérieur de la portée d'une stratégie explicite).

default_strategy = tf.distribute.get_strategy()

Cette stratégie répond à deux objectifs principaux :

  • Il permet d'écrire inconditionnellement du code de bibliothèque compatible avec la distribution. Par exemple, dans tf.optimizer s, vous pouvez utiliser tf.distribute.get_strategy et utiliser cette stratégie pour réduire les gradients. Il renverra toujours un objet de stratégie sur lequel vous pourrez appeler l'API Strategy.reduce .
# In optimizer or other library code
# Get currently active strategy
strategy = tf.distribute.get_strategy()
strategy.reduce("SUM", 1., axis=None)  # reduce some values
1.0
  • Semblable au code de bibliothèque, il peut être utilisé pour écrire des programmes d'utilisateurs finaux pour travailler avec et sans stratégie de distribution, sans nécessiter de logique conditionnelle. Voici un exemple d'extrait de code illustrant ceci :
if tf.config.list_physical_devices('GPU'):
  strategy = tf.distribute.MirroredStrategy()
else:  # Use the Default Strategy
  strategy = tf.distribute.get_strategy()

with strategy.scope():
  # Do something interesting
  print(tf.Variable(1.))
INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:GPU:0',)
MirroredVariable:{
  0: <tf.Variable 'Variable:0' shape=() dtype=float32, numpy=1.0>
}

OneDeviceStrategy

tf.distribute.OneDeviceStrategy est une stratégie pour placer toutes les variables et tous les calculs sur un seul appareil spécifié.

strategy = tf.distribute.OneDeviceStrategy(device="/gpu:0")

Cette stratégie se distingue de la stratégie par défaut de plusieurs façons. Dans la stratégie par défaut, la logique de placement des variables reste inchangée par rapport à l'exécution de TensorFlow sans aucune stratégie de distribution. Mais lors de l'utilisation OneDeviceStrategy , toutes les variables créées dans sa portée sont explicitement placées sur le périphérique spécifié. De plus, toutes les fonctions appelées via OneDeviceStrategy.run seront également placées sur l'appareil spécifié.

L'entrée distribuée via cette stratégie sera prérécupérée sur l'appareil spécifié. Dans la stratégie par défaut, il n'y a pas de distribution d'entrée.

Semblable à la stratégie par défaut, cette stratégie peut également être utilisée pour tester votre code avant de passer à d'autres stratégies qui distribuent réellement à plusieurs appareils/machines. Cela exercera un peu plus la machinerie de stratégie de distribution que la stratégie par défaut, mais pas dans toute la mesure de l'utilisation, par exemple, de MirroredStrategy ou TPUStrategy . Si vous voulez du code qui se comporte comme s'il n'y avait pas de stratégie, utilisez la stratégie par défaut.

Jusqu'à présent, vous avez appris différentes stratégies et comment vous pouvez les instancier. Les quelques sections suivantes montrent les différentes manières dont vous pouvez les utiliser pour diffuser votre formation.

Utiliser tf.distribute.Strategy avec Keras Model.fit

tf.distribute.Strategy est intégré à tf.keras , qui est l'implémentation de TensorFlow de la spécification de l'API Keras . tf.keras est une API de haut niveau pour créer et former des modèles. En s'intégrant au backend tf.keras , vous pouvez facilement distribuer votre formation écrite dans le cadre de formation Keras à l' aide de Model.fit .

Voici ce que vous devez modifier dans votre code :

  1. Créez une instance du tf.distribute.Strategy approprié.
  2. Déplacez la création du modèle, de l'optimiseur et des métriques Keras dans strategy.scope .

Les stratégies de distribution TensorFlow prennent en charge tous les types de modèles Keras : Sequential , Functional et subclassed .

Voici un extrait de code pour le faire pour un modèle Keras très simple avec une couche Dense :

mirrored_strategy = tf.distribute.MirroredStrategy()

with mirrored_strategy.scope():
  model = tf.keras.Sequential([tf.keras.layers.Dense(1, input_shape=(1,))])

model.compile(loss='mse', optimizer='sgd')
INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:GPU:0',)
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).

Cet exemple utilise MirroredStrategy , vous pouvez donc l'exécuter sur une machine avec plusieurs GPU. strategy.scope() indique à Keras quelle stratégie utiliser pour distribuer la formation. La création de modèles/optimiseurs/métriques dans cette portée vous permet de créer des variables distribuées au lieu de variables régulières. Une fois que cela est configuré, vous pouvez adapter votre modèle comme vous le feriez normalement. MirroredStrategy se charge de répliquer la formation du modèle sur les GPU disponibles, d'agréger les gradients, etc.

dataset = tf.data.Dataset.from_tensors(([1.], [1.])).repeat(100).batch(10)
model.fit(dataset, epochs=2)
model.evaluate(dataset)
Epoch 1/2
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
2021-10-26 01:27:56.527729: W tensorflow/core/grappler/optimizers/data/auto_shard.cc:695] AUTO sharding policy will apply DATA sharding policy as it failed to apply FILE sharding policy because of the following reason: Found an unshardable source dataset: name: "TensorDataset/_2"
op: "TensorDataset"
input: "Placeholder/_0"
input: "Placeholder/_1"
attr {
  key: "Toutput_types"
  value {
    list {
      type: DT_FLOAT
      type: DT_FLOAT
    }
  }
}
attr {
  key: "output_shapes"
  value {
    list {
      shape {
        dim {
          size: 1
        }
      }
      shape {
        dim {
          size: 1
        }
      }
    }
  }
}
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
10/10 [==============================] - 3s 2ms/step - loss: 2.2552
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
Epoch 2/2
10/10 [==============================] - 0s 2ms/step - loss: 0.9968
2021-10-26 01:27:59.372113: W tensorflow/core/grappler/optimizers/data/auto_shard.cc:695] AUTO sharding policy will apply DATA sharding policy as it failed to apply FILE sharding policy because of the following reason: Found an unshardable source dataset: name: "TensorDataset/_2"
op: "TensorDataset"
input: "Placeholder/_0"
input: "Placeholder/_1"
attr {
  key: "Toutput_types"
  value {
    list {
      type: DT_FLOAT
      type: DT_FLOAT
    }
  }
}
attr {
  key: "output_shapes"
  value {
    list {
      shape {
        dim {
          size: 1
        }
      }
      shape {
        dim {
          size: 1
        }
      }
    }
  }
}
10/10 [==============================] - 1s 2ms/step - loss: 0.6190
0.6190494298934937

Ici, un tf.data.Dataset fournit l'entrée de formation et d'évaluation. Vous pouvez également utiliser des tableaux NumPy :

import numpy as np

inputs, targets = np.ones((100, 1)), np.ones((100, 1))
model.fit(inputs, targets, epochs=2, batch_size=10)
Epoch 1/2
2021-10-26 01:28:00.609977: W tensorflow/core/grappler/optimizers/data/auto_shard.cc:695] AUTO sharding policy will apply DATA sharding policy as it failed to apply FILE sharding policy because of the following reason: Did not find a shardable source, walked to a node which is not a dataset: name: "FlatMapDataset/_9"
op: "FlatMapDataset"
input: "PrefetchDataset/_8"
attr {
  key: "Targuments"
  value {
    list {
    }
  }
}
attr {
  key: "f"
  value {
    func {
      name: "__inference_Dataset_flat_map_slice_batch_indices_997"
    }
  }
}
attr {
  key: "output_shapes"
  value {
    list {
      shape {
        dim {
          size: 10
        }
      }
    }
  }
}
attr {
  key: "output_types"
  value {
    list {
      type: DT_INT64
    }
  }
}
. Consider either turning off auto-sharding or switching the auto_shard_policy to DATA to shard this dataset. You can do this by creating a new `tf.data.Options()` object then setting `options.experimental_distribute.auto_shard_policy = AutoShardPolicy.DATA` before applying the options object to the dataset via `dataset.with_options(options)`.
10/10 [==============================] - 1s 2ms/step - loss: 0.4406
Epoch 2/2
10/10 [==============================] - 0s 2ms/step - loss: 0.1947
<keras.callbacks.History at 0x7fb81813d2d0>

Dans les deux cas, avec Dataset ou NumPy, chaque lot de l'entrée donnée est divisé également entre les multiples répliques. Par exemple, si vous utilisez MirroredStrategy avec 2 GPU, chaque lot de taille 10 sera divisé entre les 2 GPU, chacun recevant 5 exemples d'entrée à chaque étape. Chaque époque s'entraînera alors plus rapidement à mesure que vous ajouterez plus de GPU. En règle générale, vous souhaiterez augmenter la taille de votre lot à mesure que vous ajoutez d'autres accélérateurs, afin d'utiliser efficacement la puissance de calcul supplémentaire. Vous devrez également réajuster votre taux d'apprentissage, selon le modèle. Vous pouvez utiliser strategy.num_replicas_in_sync pour obtenir le nombre de répliques.

# Compute a global batch size using a number of replicas.
BATCH_SIZE_PER_REPLICA = 5
global_batch_size = (BATCH_SIZE_PER_REPLICA *
                     mirrored_strategy.num_replicas_in_sync)
dataset = tf.data.Dataset.from_tensors(([1.], [1.])).repeat(100)
dataset = dataset.batch(global_batch_size)

LEARNING_RATES_BY_BATCH_SIZE = {5: 0.1, 10: 0.15}
learning_rate = LEARNING_RATES_BY_BATCH_SIZE[global_batch_size]

Qu'est-ce qui est pris en charge maintenant ?

API de formation MirroredStrategy TPUStrategy MultiWorkerMirroredStrategy ParameterServerStrategy CentralStorageStrategy
Keras Model.fit Prise en charge Prise en charge Prise en charge Support expérimental Support expérimental

Exemples et tutoriels

Voici une liste de tutoriels et d'exemples qui illustrent l'intégration ci-dessus de bout en bout avec Keras Model.fit :

  1. Tutoriel : Entraînement avec Model.fit et MirroredStrategy .
  2. Tutoriel : Entraînement avec Model.fit et MultiWorkerMirroredStrategy .
  3. Guide : contient un exemple d'utilisation de Model.fit et TPUStrategy .
  4. Tutoriel : Entraînement du serveur de paramètres avec Model.fit et ParameterServerStrategy .
  5. Tutoriel : affiner le BERT pour de nombreuses tâches du benchmark GLUE avec Model.fit et TPUStrategy .
  6. Référentiel TensorFlow Model Garden contenant des collections de modèles de pointe mis en œuvre à l'aide de diverses stratégies.

Utiliser tf.distribute.Strategy avec des boucles d'entraînement personnalisées

Comme démontré ci-dessus, l'utilisation de tf.distribute.Strategy avec Keras Model.fit nécessite de modifier seulement quelques lignes de votre code. Avec un peu plus d'effort, vous pouvez également utiliser tf.distribute.Strategy avec des boucles d'entraînement personnalisées .

Si vous avez besoin de plus de flexibilité et de contrôle sur vos boucles d'entraînement qu'avec Estimator ou Keras, vous pouvez écrire des boucles d'entraînement personnalisées. Par exemple, lorsque vous utilisez un GAN, vous voudrez peut-être effectuer un nombre différent de pas de générateur ou de discriminateur à chaque tour. De même, les cadres de haut niveau sont peu adaptés à la formation en Apprentissage par Renforcement.

Les classes tf.distribute.Strategy fournissent un ensemble de méthodes de base pour prendre en charge les boucles de formation personnalisées. Leur utilisation peut nécessiter une restructuration mineure du code au départ, mais une fois cela fait, vous devriez pouvoir basculer entre les GPU, les TPU et plusieurs machines simplement en changeant l'instance de stratégie.

Vous trouverez ci-dessous un bref extrait illustrant ce cas d'utilisation pour un exemple de formation simple utilisant le même modèle Keras qu'auparavant.

Tout d'abord, créez le modèle et l'optimiseur dans le cadre de la stratégie. Cela garantit que toutes les variables créées avec le modèle et l'optimiseur sont des variables en miroir.

with mirrored_strategy.scope():
  model = tf.keras.Sequential([tf.keras.layers.Dense(1, input_shape=(1,))])
  optimizer = tf.keras.optimizers.SGD()

Ensuite, créez l'ensemble de données d'entrée et appelez tf.distribute.Strategy.experimental_distribute_dataset pour distribuer l'ensemble de données en fonction de la stratégie.

dataset = tf.data.Dataset.from_tensors(([1.], [1.])).repeat(100).batch(
    global_batch_size)
dist_dataset = mirrored_strategy.experimental_distribute_dataset(dataset)
2021-10-26 01:28:01.831942: W tensorflow/core/grappler/optimizers/data/auto_shard.cc:695] AUTO sharding policy will apply DATA sharding policy as it failed to apply FILE sharding policy because of the following reason: Found an unshardable source dataset: name: "TensorDataset/_2"
op: "TensorDataset"
input: "Placeholder/_0"
input: "Placeholder/_1"
attr {
  key: "Toutput_types"
  value {
    list {
      type: DT_FLOAT
      type: DT_FLOAT
    }
  }
}
attr {
  key: "output_shapes"
  value {
    list {
      shape {
        dim {
          size: 1
        }
      }
      shape {
        dim {
          size: 1
        }
      }
    }
  }
}

Ensuite, définissez une étape de la formation. Utilisez tf.GradientTape pour calculer les gradients et l'optimiseur pour appliquer ces gradients afin de mettre à jour les variables de votre modèle. Pour distribuer cette étape de formation, placez-la dans une fonction train_step et transmettez-la à tf.distribute.Strategy.run avec les entrées de l'ensemble de données que vous avez obtenues à partir du dist_dataset créé auparavant :

loss_object = tf.keras.losses.BinaryCrossentropy(
  from_logits=True,
  reduction=tf.keras.losses.Reduction.NONE)

def compute_loss(labels, predictions):
  per_example_loss = loss_object(labels, predictions)
  return tf.nn.compute_average_loss(per_example_loss, global_batch_size=global_batch_size)

def train_step(inputs):
  features, labels = inputs

  with tf.GradientTape() as tape:
    predictions = model(features, training=True)
    loss = compute_loss(labels, predictions)

  gradients = tape.gradient(loss, model.trainable_variables)
  optimizer.apply_gradients(zip(gradients, model.trainable_variables))
  return loss

@tf.function
def distributed_train_step(dist_inputs):
  per_replica_losses = mirrored_strategy.run(train_step, args=(dist_inputs,))
  return mirrored_strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses,
                         axis=None)

Quelques autres choses à noter dans le code ci-dessus :

  1. Vous avez utilisé tf.nn.compute_average_loss pour calculer la perte. tf.nn.compute_average_loss la perte par exemple et divise la somme par global_batch_size . Ceci est important car plus tard, après le calcul des gradients sur chaque réplique, ils sont agrégés sur les répliques en les additionnant .
  2. Vous avez également utilisé l'API tf.distribute.Strategy.reduce pour agréger les résultats renvoyés par tf.distribute.Strategy.run . tf.distribute.Strategy.run renvoie les résultats de chaque réplica local de la stratégie, et il existe plusieurs façons d'utiliser ce résultat. Vous pouvez les reduce pour obtenir une valeur agrégée. Vous pouvez également faire tf.distribute.Strategy.experimental_local_results pour obtenir la liste des valeurs contenues dans le résultat, une par réplica local.
  3. Lorsque vous appelez apply_gradients dans une portée de stratégie de distribution, son comportement est modifié. Plus précisément, avant d'appliquer des gradients sur chaque instance parallèle lors de l'apprentissage synchrone, il effectue une somme sur toutes les répliques des gradients.

Enfin, une fois que vous avez défini l'étape d'entraînement, vous pouvez parcourir dist_dataset et exécuter l'entraînement en boucle :

for dist_inputs in dist_dataset:
  print(distributed_train_step(dist_inputs))
tf.Tensor(0.18686396, shape=(), dtype=float32)
tf.Tensor(0.18628375, shape=(), dtype=float32)
tf.Tensor(0.18570684, shape=(), dtype=float32)
tf.Tensor(0.18513316, shape=(), dtype=float32)
tf.Tensor(0.1845627, shape=(), dtype=float32)
tf.Tensor(0.18399543, shape=(), dtype=float32)
tf.Tensor(0.18343134, shape=(), dtype=float32)
tf.Tensor(0.18287037, shape=(), dtype=float32)
tf.Tensor(0.18231256, shape=(), dtype=float32)
tf.Tensor(0.18175781, shape=(), dtype=float32)
tf.Tensor(0.18120615, shape=(), dtype=float32)
tf.Tensor(0.18065754, shape=(), dtype=float32)
tf.Tensor(0.18011193, shape=(), dtype=float32)
tf.Tensor(0.17956935, shape=(), dtype=float32)
tf.Tensor(0.17902976, shape=(), dtype=float32)
tf.Tensor(0.17849308, shape=(), dtype=float32)
tf.Tensor(0.17795937, shape=(), dtype=float32)
tf.Tensor(0.17742859, shape=(), dtype=float32)
tf.Tensor(0.17690066, shape=(), dtype=float32)
tf.Tensor(0.17637561, shape=(), dtype=float32)

Dans l'exemple ci-dessus, vous avez itéré sur le dist_dataset pour fournir une entrée à votre formation. Vous disposez également de l'ensemble de tf.distribute.Strategy.make_experimental_numpy_dataset pour prendre en charge les entrées NumPy. Vous pouvez utiliser cette API pour créer un ensemble de données avant d'appeler tf.distribute.Strategy.experimental_distribute_dataset .

Une autre façon d'itérer sur vos données consiste à utiliser explicitement des itérateurs. Vous souhaiterez peut-être effectuer cette opération lorsque vous souhaitez exécuter un nombre donné d'étapes au lieu d'itérer sur l'ensemble du jeu de données. L'itération ci-dessus serait maintenant modifiée pour créer d'abord un itérateur, puis appeler explicitement next pour obtenir les données d'entrée.

iterator = iter(dist_dataset)
for _ in range(10):
  print(distributed_train_step(next(iterator)))
tf.Tensor(0.17585339, shape=(), dtype=float32)
tf.Tensor(0.17533402, shape=(), dtype=float32)
tf.Tensor(0.17481743, shape=(), dtype=float32)
tf.Tensor(0.17430364, shape=(), dtype=float32)
tf.Tensor(0.17379259, shape=(), dtype=float32)
tf.Tensor(0.17328428, shape=(), dtype=float32)
tf.Tensor(0.17277871, shape=(), dtype=float32)
tf.Tensor(0.17227581, shape=(), dtype=float32)
tf.Tensor(0.17177561, shape=(), dtype=float32)
tf.Tensor(0.17127804, shape=(), dtype=float32)

Cela couvre le cas le plus simple d'utilisation de l'API tf.distribute.Strategy pour distribuer des boucles de formation personnalisées.

Qu'est-ce qui est pris en charge maintenant ?

API de formation MirroredStrategy TPUStrategy MultiWorkerMirroredStrategy ParameterServerStrategy CentralStorageStrategy
Boucle d'entraînement personnalisée Prise en charge Prise en charge Prise en charge Support expérimental Support expérimental

Exemples et tutoriels

Voici quelques exemples d'utilisation de stratégies de distribution avec des boucles d'entraînement personnalisées :

  1. Tutoriel : Entraînement avec une boucle d'entraînement personnalisée et MirroredStrategy .
  2. Tutoriel : Entraînement avec une boucle d'entraînement personnalisée et MultiWorkerMirroredStrategy .
  3. Guide : contient un exemple de boucle d'entraînement personnalisée avec TPUStrategy .
  4. Tutoriel : Entraînement du serveur de paramètres avec une boucle d'entraînement personnalisée et ParameterServerStrategy .
  5. Référentiel TensorFlow Model Garden contenant des collections de modèles de pointe mis en œuvre à l'aide de diverses stratégies.

Autres sujets

Cette section couvre certains sujets pertinents pour plusieurs cas d'utilisation.

Paramétrage de la variable d'environnement TF_CONFIG

Pour la formation multi-travailleur, comme mentionné précédemment, vous devez configurer la variable d'environnement 'TF_CONFIG' pour chaque binaire exécuté dans votre cluster. La variable d'environnement 'TF_CONFIG' est une chaîne JSON qui spécifie quelles tâches constituent un cluster, leurs adresses et le rôle de chaque tâche dans le cluster. Le tensorflow/ecosystem fournit un modèle Kubernetes, qui configure 'TF_CONFIG' pour vos tâches de formation.

Il y a deux composants de 'TF_CONFIG' : un cluster et une tâche.

  • Un cluster fournit des informations sur le cluster de formation, qui est un dict composé de différents types d'emplois tels que les travailleurs. Dans la formation multi-travailleurs, il y a généralement un travailleur qui assume un peu plus de responsabilités, comme enregistrer un point de contrôle et écrire un fichier récapitulatif pour TensorBoard en plus de ce que fait un travailleur régulier. Un tel travailleur est appelé le travailleur "chef", et il est d'usage que le travailleur avec l'indice 0 soit nommé travailleur principal (en fait, c'est ainsi que tf.distribute.Strategy est implémenté).
  • Une tâche, quant à elle, fournit des informations sur la tâche en cours. Le premier cluster de composants est le même pour tous les travailleurs, et la deuxième tâche de composant est différente sur chaque travailleur et spécifie le type et l'index de ce travailleur.

Un exemple de 'TF_CONFIG' est :

os.environ["TF_CONFIG"] = json.dumps({
    "cluster": {
        "worker": ["host1:port", "host2:port", "host3:port"],
        "ps": ["host4:port", "host5:port"]
    },
   "task": {"type": "worker", "index": 1}
})

Ce 'TF_CONFIG' spécifie qu'il y a trois travailleurs et deux tâches "ps" dans le "cluster" avec leurs hôtes et leurs ports. La "task" précise le rôle de la tâche en cours dans le "cluster" —worker 1 (le second worker). Les rôles valides dans un cluster sont "chief" , "worker" , "ps" et "evaluator" . Il ne devrait pas y avoir de travail "ps" sauf lors de l'utilisation tf.distribute.experimental.ParameterServerStrategy .

Et après?

tf.distribute.Strategy est activement en cours de développement. Essayez-le et fournissez vos commentaires en utilisant les problèmes GitHub .