Vertex AI: entraînement multinœud et apprentissage par transfert avec TensorFlow

1. Présentation

Dans cet atelier, vous allez utiliser Vertex AI afin d'exécuter un job d'entraînement à plusieurs nœuds de calcul pour un modèle TensorFlow.

Objectifs

Vous allez apprendre à effectuer les opérations suivantes :

  • Modifier le code de l'application d'entraînement pour l'entraînement à plusieurs nœuds de calcul
  • Configurer et lancer un job d'entraînement à plusieurs nœuds depuis l'interface utilisateur de Vertex AI
  • Configurer et lancer un job d'entraînement à plusieurs nœuds avec Vertex SDK

Le coût total d'exécution de cet atelier sur Google Cloud est d'environ 5$.

2. Présentation de Vertex AI

Cet atelier utilise la toute dernière offre de produits d'IA de Google Cloud. Vertex AI simplifie l'expérience de développement en intégrant toutes les offres de ML de Google Cloud. Auparavant, les modèles entraînés avec AutoML et les modèles personnalisés étaient accessibles depuis des services distincts. La nouvelle offre regroupe ces deux types de modèles mais aussi d'autres nouveaux produits en une seule API. Vous pouvez également migrer des projets existants vers Vertex AI. Pour envoyer un commentaire, consultez la page d'assistance.

Vertex AI comprend de nombreux produits différents qui permettent de gérer les workflows de ML de bout en bout. Cet atelier se concentre sur les produits mis en évidence ci-dessous : Training et Workbench.

Présentation des produits Vertex

3. Présentation du cas d'utilisation

Dans cet atelier, vous allez utiliser l'apprentissage par transfert pour entraîner un modèle de classification d'images sur l'ensemble de données cassava à partir d'ensembles de données TensorFlow. L'architecture que vous allez utiliser est un modèle ResNet50 de la bibliothèque tf.keras.applications pré-entraîné sur l'ensemble de données Imagenet.

Pourquoi utiliser l'entraînement distribué ?

Si vous disposez d'un seul GPU, TensorFlow utilisera cet accélérateur pour entraîner le modèle plus rapidement, sans qu'aucune action supplémentaire de votre part ne soit nécessaire. Toutefois, si vous souhaitez profiter d'un avantage supplémentaire en utilisant plusieurs GPU sur une même machine ou sur plusieurs machines (chacune possédant potentiellement plusieurs GPU), vous devez utiliser tf.distribute, la bibliothèque de TensorFlow permettant d'exécuter des calculs sur plusieurs appareils. Le terme "appareil" désigne un processeur ou un accélérateur, tel que des GPU ou des TPU, qui est installé sur une machine sur laquelle TensorFlow peut exécuter des opérations.

Le moyen le plus simple de commencer à utiliser l'entraînement distribué est d'utiliser une seule machine avec plusieurs appareils GPU. Une stratégie de distribution TensorFlow du module tf.distribute gère la coordination de la distribution des données et les mises à jour des gradients sur tous les GPU. Si vous maîtrisez l'entraînement sur un seul hôte et que vous souhaitez évoluer encore plus, ajouter plusieurs machines à votre cluster peut vous aider à améliorer encore davantage les performances. Vous pouvez utiliser un cluster de machines disposant uniquement d'un processeur, ou possédant chacune un ou plusieurs GPU. Dans cet atelier, vous allez découvrir comment utiliser MultiWorkerMirroredStrategy pour distribuer l'entraînement d'un modèle TensorFlow sur plusieurs machines sur Vertex AI.

MultiWorkerMirroredStrategy est une stratégie de parallélisme des données synchrone que vous pouvez utiliser avec seulement quelques modifications de code. Une copie du modèle est créée sur chaque appareil de votre cluster. Les mises à jour de gradient suivantes s'effectueront de manière synchrone. Cela signifie que chaque appareil de calcul calcule la propagation avant et arrière du modèle sur une tranche différente des données d'entrée. Les gradients calculés à partir de chacune de ces tranches sont ensuite agrégés sur l'ensemble des appareils d'une machine et de toutes les machines du cluster, puis réduits (généralement une moyenne) au cours d'un processus appelé "all-reduce". L'optimiseur met ensuite à jour les paramètres avec ces gradients réduits, ce qui permet de synchroniser les appareils. Pour en savoir plus sur l'entraînement distribué avec TensorFlow, regardez la vidéo ci-dessous:

4. Configurer votre environnement

Pour suivre cet atelier de programmation, vous aurez besoin d'un projet Google Cloud Platform dans lequel la facturation est activée. Pour créer un projet, suivez ces instructions.

Étape 1 : Activez l'API Compute Engine

Accédez à Compute Engine et cliquez sur Activer si ce n'est pas déjà fait. Vous en aurez besoin pour créer votre instance de notebook.

Étape 2 : Activez l'API Container Registry

Accédez à Container Registry et cliquez sur Activer si ce n'est pas déjà fait. Vous utiliserez ce service afin de créer un conteneur pour votre tâche d'entraînement personnalisé.

Étape 3 : Activez l'API Vertex AI

Accédez à la section Vertex AI de Cloud Console, puis cliquez sur Activer l'API Vertex AI.

Tableau de bord Vertex AI

Étape 4 : Créez une instance Vertex AI Workbench

Dans la section Vertex AI de Cloud Console, cliquez sur Workbench :

Menu Vertex AI

Activez l'API Notebooks si ce n'est pas déjà fait.

Notebook_api

Une fois l'API activée, cliquez sur NOTEBOOKS GÉRÉS :

Notebooks_UI

Sélectionnez ensuite NOUVEAU NOTEBOOK.

new_notebook

Attribuez un nom à votre notebook, puis cliquez sur Paramètres avancés.

create_notebook

Sous "Paramètres avancés", activez l'arrêt en cas d'inactivité et définissez le nombre de minutes sur 60. Cela entraîne l'arrêt automatique du notebook lorsqu'il n'est pas utilisé. Vous ne payez donc pas de frais inutiles.

idle_timeout

Dans la section Sécurité, sélectionnez "Activer le terminal" si ce n'est pas déjà fait.

enable_terminal

Vous pouvez conserver tous les autres paramètres avancés tels quels.

Cliquez ensuite sur Créer. Le provisionnement de l'instance prend quelques minutes.

Une fois l'instance créée, sélectionnez Ouvrir JupyterLab.

open_jupyterlab

La première fois que vous utilisez une nouvelle instance, vous êtes invité à vous authentifier. Pour ce faire, suivez les étapes indiquées dans l'interface utilisateur.

authenticate

5. Conteneuriser le code de l'application d'entraînement

Vous allez envoyer ce job d'entraînement à Vertex en plaçant le code de votre application d'entraînement dans un conteneur Docker et en déployant ce conteneur dans Google Container Registry. Cette approche vous permet d'entraîner un modèle créé avec n'importe quel framework.

Pour commencer, à partir du menu de lancement, ouvrez une fenêtre de terminal dans votre instance de notebook :

Ouvrir le terminal dans le notebook

Créez un répertoire appelé cassava et utilisez la commande cd pour y accéder :

mkdir cassava
cd cassava

Étape 1: Créez un Dockerfile

La première étape de la conteneurisation de votre code consiste à créer un Dockerfile. Vous allez placer dans ce Dockerfile toutes les commandes nécessaires à l'exécution de l'image. Il installera toutes les bibliothèques nécessaires et configurera le point d'entrée du code d'entraînement.

Depuis votre terminal, créez un Dockerfile vide:

touch Dockerfile

Ouvrez le Dockerfile et copiez-y le code suivant :

FROM gcr.io/deeplearning-platform-release/tf2-gpu.2-7

WORKDIR /

# Copies the trainer code to the docker image.
COPY trainer /trainer

# Sets up the entry point to invoke the trainer.
ENTRYPOINT ["python", "-m", "trainer.task"]

Ce fichier Dockerfile utilise l'image Docker avec GPU de TensorFlow Enterprise 2.7 comme conteneur de deep learning. Les conteneurs de deep learning sur Google Cloud sont fournis avec de nombreux frameworks de ML et de data science courants préinstallés. Une fois cette image téléchargée, ce Dockerfile configure le point d'entrée du code d'entraînement. Vous n'avez pas encore créé ces fichiers. À la prochaine étape, vous allez ajouter le code d'entraînement et de réglage du modèle.

Étape 2: Créez un bucket Cloud Storage

Dans ce job d'entraînement, vous allez exporter le modèle TensorFlow entraîné vers un bucket Cloud Storage. Depuis votre terminal, exécutez la commande suivante afin de définir une variable d'environnement pour votre projet, en veillant à remplacer your-cloud-project par l'ID de votre projet :

PROJECT_ID='your-cloud-project'

Ensuite, exécutez la commande suivante dans votre terminal pour créer un bucket dans votre projet.

BUCKET="gs://${PROJECT_ID}-bucket"
gsutil mb -l us-central1 $BUCKET

Étape 3: Ajoutez le code d'entraînement du modèle

Depuis votre terminal, exécutez le code suivant afin de créer un répertoire pour le code d'entraînement et un fichier Python dans lequel vous ajouterez le code :

mkdir trainer
touch trainer/task.py

Votre répertoire cassava/ doit maintenant contenir les éléments suivants :

+ Dockerfile
+ trainer/
    + task.py

Ouvrez ensuite le fichier task.py que vous venez de créer et copiez-y le code ci-dessous. Vous devez remplacer {your-gcs-bucket} par le nom du bucket Cloud Storage que vous venez de créer.

import tensorflow as tf
import tensorflow_datasets as tfds
import os


PER_REPLICA_BATCH_SIZE = 64
EPOCHS = 2

# TODO: replace {your-gcs-bucket} with the name of the Storage bucket you created earlier
BUCKET = 'gs://{your-gcs-bucket}/mwms'

def preprocess_data(image, label):
  '''Resizes and scales images.'''

  image = tf.image.resize(image, (300,300))
  return tf.cast(image, tf.float32) / 255., label


def create_dataset(batch_size):
  '''Loads Cassava dataset and preprocesses data.'''

  data, info = tfds.load(name='cassava', as_supervised=True, with_info=True)
  number_of_classes = info.features['label'].num_classes
  train_data = data['train'].map(preprocess_data,
                                 num_parallel_calls=tf.data.experimental.AUTOTUNE)
  train_data  = train_data.shuffle(1000)
  train_data  = train_data.batch(batch_size)
  train_data  = train_data.prefetch(tf.data.experimental.AUTOTUNE)

  # Set AutoShardPolicy
  options = tf.data.Options()
  options.experimental_distribute.auto_shard_policy = tf.data.experimental.AutoShardPolicy.DATA
  train_data = train_data.with_options(options)

  return train_data, number_of_classes


def create_model(number_of_classes):
  '''Creates and compiles pretrained ResNet50 model.'''

  base_model = tf.keras.applications.ResNet50(weights='imagenet', include_top=False)
  x = base_model.output
  x = tf.keras.layers.GlobalAveragePooling2D()(x)
  x = tf.keras.layers.Dense(1016, activation='relu')(x)
  predictions = tf.keras.layers.Dense(number_of_classes, activation='softmax')(x)
  model = tf.keras.Model(inputs=base_model.input, outputs=predictions)

  model.compile(
      loss='sparse_categorical_crossentropy',
      optimizer=tf.keras.optimizers.Adam(0.0001),
      metrics=['accuracy'])

  return model


def _is_chief(task_type, task_id):
  '''Helper function. Determines if machine is chief.'''

  return task_type == 'chief'


def _get_temp_dir(dirpath, task_id):
  '''Helper function. Gets temporary directory for saving model.'''

  base_dirpath = 'workertemp_' + str(task_id)
  temp_dir = os.path.join(dirpath, base_dirpath)
  tf.io.gfile.makedirs(temp_dir)
  return temp_dir


def write_filepath(filepath, task_type, task_id):
  '''Helper function. Gets filepath to save model.'''

  dirpath = os.path.dirname(filepath)
  base = os.path.basename(filepath)
  if not _is_chief(task_type, task_id):
    dirpath = _get_temp_dir(dirpath, task_id)
  return os.path.join(dirpath, base)


def main():
  # Create strategy
  strategy = tf.distribute.MultiWorkerMirroredStrategy()

  # Get data
  global_batch_size = PER_REPLICA_BATCH_SIZE * strategy.num_replicas_in_sync
  train_data, number_of_classes = create_dataset(global_batch_size)

  # Wrap variable creation within strategy scope
  with strategy.scope():
    model = create_model(number_of_classes)

  model.fit(train_data, epochs=EPOCHS)

  # Determine type and task of the machine from
  # the strategy cluster resolver
  task_type, task_id = (strategy.cluster_resolver.task_type,
                        strategy.cluster_resolver.task_id)

  # Based on the type and task, write to the desired model path
  write_model_path = write_filepath(BUCKET, task_type, task_id)
  model.save(write_model_path)

if __name__ == "__main__":
    main()

Avant de créer le conteneur, examinons le code qui utilise MultiWorkerMirroredStrategy de l'API tf.distribute.Strategy.

Certains composants du code sont nécessaires pour qu'il fonctionne avec MultiWorkerMirroredStrategy.

  1. Les données doivent être segmentées, ce qui signifie que chaque nœud de calcul se voit attribuer un sous-ensemble de l'ensemble de données complet. Par conséquent, à chaque étape, une taille de lot globale d'éléments d'ensemble de données qui ne se chevauchent pas sera traitée par chaque nœud de calcul. Cette segmentation s'effectue automatiquement avec tf.data.experimental.AutoShardPolicy, qui peut être défini sur FILE ou DATA. Dans cet exemple, la fonction create_dataset() définit AutoShardPolicy sur DATA, car le jeu de données sur le manioc n'est pas téléchargé sous forme de plusieurs fichiers. Toutefois, si vous ne définissez pas la règle sur DATA, la règle AUTO par défaut est appliquée et le résultat final est le même. Pour en savoir plus sur la segmentation des ensembles de données avec MultiWorkerMirroredStrategy, cliquez ici.
  2. L'objet MultiWorkerMirroredStrategy est créé dans la fonction main(). Vous devez ensuite encapsuler la création des variables du modèle dans le champ d'application de la stratégie. Cette étape essentielle indique à TensorFlow les variables à mettre en miroir sur les instances répliquées.
  3. La taille du lot est mise à l'échelle par la valeur num_replicas_in_sync. Cela garantit que chaque instance répliquée traite le même nombre d'exemples à chaque étape. La mise à l'échelle de la taille de lot est une bonne pratique lorsque vous utilisez des stratégies de chargement de données synchrones en parallèle dans TensorFlow.
  4. L'enregistrement de votre modèle est un peu plus compliqué avec plusieurs nœuds de calcul, car la destination doit être différente pour chacun des nœuds de calcul. Le nœud de calcul chef procède à l'enregistrement dans le répertoire du modèle souhaité, tandis que les autres nœuds de calcul enregistrent le modèle dans des répertoires temporaires. Ces répertoires temporaires doivent impérativement être uniques afin d'éviter que plusieurs nœuds écrivent au même emplacement. Les opérations d'enregistrement peuvent être collectives. En d'autres termes, tous les nœuds doivent procéder à l'enregistrement et pas seulement le nœud chef. Les fonctions _is_chief(), _get_temp_dir(), write_filepath() et la fonction main() incluent toutes du code récurrent qui permet de sauvegarder le modèle.

Remarque : si vous avez utilisé MultiWorkerMirroredStrategy dans un environnement différent, vous avez peut-être configuré la variable d'environnement TF_CONFIG. Vertex AI définit automatiquement TF_CONFIG. Vous n'avez donc pas besoin de définir cette variable sur chaque machine de votre cluster.

Étape 4 : Créez le conteneur

Depuis votre terminal, exécutez la commande suivante afin de définir une variable d'environnement pour votre projet, en veillant à remplacer your-cloud-project par l'ID de votre projet :

PROJECT_ID='your-cloud-project'

Définissez une variable avec l'URI de votre image de conteneur dans Google Container Registry :

IMAGE_URI="gcr.io/$PROJECT_ID/multiworker:cassava"

Configurez Docker.

gcloud auth configure-docker

Ensuite, créez le conteneur en exécutant la commande suivante à partir de la racine de votre répertoire cassava :

docker build ./ -t $IMAGE_URI

Enfin, transférez-le vers Google Container Registry :

docker push $IMAGE_URI

Maintenant que le conteneur a été transféré vers Container Registry, vous êtes prêt à lancer un job d'entraînement.

6. Exécuter un job d'entraînement à plusieurs nœuds sur Vertex AI

Cet atelier utilise un entraînement personnalisé via un conteneur personnalisé de Google Container Registry, mais vous pouvez également exécuter un job d'entraînement avec les conteneurs prédéfinis.

Pour commencer, accédez à l'onglet Entraînement de la section Vertex de votre console Cloud :

Menu uCAIP

Étape 1 : Configurez le job d'entraînement

Cliquez sur Créer pour saisir les paramètres de votre job d'entraînement.

  • Sous Ensemble de données, sélectionnez Aucun ensemble de données géré.
  • Sélectionnez ensuite Entraînement personnalisé (avancé) comme méthode d'entraînement, puis cliquez sur Continuer.
  • Saisissez multiworker-cassava (ou tout autre nom que vous souhaitez donner à votre modèle) dans Nom du modèle.
  • Cliquez sur Continuer.

À l'étape des paramètres du conteneur, sélectionnez Conteneur personnalisé :

Option de conteneur personnalisé

Dans le premier champ (Image du conteneur), saisissez la valeur de votre variable IMAGE_URI de la section précédente. Il doit s'agir de gcr.io/your-cloud-project/multiworker:cassava, avec l'ID de votre propre projet. Laissez les autres champs vides et cliquez sur Continuer.

Ignorez l'étape concernant les hyperparamètres en cliquant à nouveau sur Continuer.

Étape 2 : Configurez le cluster de calcul

Vertex AI fournit quatre pools de nœuds de calcul pour couvrir les différents types de tâches de la machine.

Le pool de nœuds de calcul 0 configure le nœud principal, chef, programmeur ou "maître". Dans MultiWorkerMirroredStrategy, toutes les machines sont désignées comme des nœuds de calcul, c'est-à-dire les machines physiques sur lesquelles le calcul répliqué est exécuté. Outre le fait que chaque machine constitue un nœud de calcul, l'un des ces nœuds doit effectuer des tâches supplémentaires, telles que l'enregistrement de points de contrôle et l'écriture de fichiers récapitulatifs sur TensorBoard. Cette machine est appelée "chef". Il n'y a qu'un nœud de calcul chef. Le nombre de vos nœuds de calcul pour le pool de nœuds de calcul 0 sera donc toujours 1.

Dans Options de calcul et tarifs, ne modifiez pas la région sélectionnée et configurez le pool de nœuds de calcul 0 de la manière suivante :

Worker_pool_0

Le pool de nœuds de calcul 1 est l'endroit où vous configurez les nœuds de calcul de votre cluster.

Configurez le pool de nœuds de calcul 1 comme suit:

Worker_pool_1

Le cluster est maintenant configuré pour disposer de deux machines à processeur uniquement. Lorsque le code de l'application d'entraînement est exécuté, MultiWorkerMirroredStrategy distribue l'entraînement sur les deux machines.

MultiWorkerMirroredStrategy ne dispose que des tâches principales et de nœud de calcul. Il n'est donc pas nécessaire de configurer les pools de nœuds de calcul supplémentaires. Toutefois, si vous utilisez le ParameterServerStrategy de TensorFlow, vous devez configurer vos serveurs de paramètres dans le pool de nœuds de calcul 2. Si vous souhaitez ajouter un évaluateur à votre cluster, vous devez configurer cette machine dans le pool de nœuds de calcul 3.

Cliquez sur Démarrer l'entraînement pour lancer le job de réglage des hyperparamètres. Votre nouveau job s'affiche dans l'onglet PIPELINES D'ENTRAÎNEMENT de la section "Entraînement" de votre console :

Jobs d'entraînement

🎉 Félicitations ! 🎉

Vous savez désormais utiliser Vertex AI pour :

  • Lancer un job d'entraînement à plusieurs nœuds pour le code d'entraînement fourni dans un conteneur personnalisé Vous avez utilisé un modèle TensorFlow dans cet exemple, mais vous pouvez entraîner un modèle créé avec n'importe quel framework à l'aide de conteneurs personnalisés ou prédéfinis.

Pour en savoir plus sur les différents composants de Vertex, consultez la documentation.

7. [Facultatif] Utiliser Vertex SDK

Dans la section précédente, vous avez découvert comment lancer le job d'entraînement via l'interface utilisateur. Dans cette section, vous allez voir une autre façon de soumettre le job d'entraînement à l'aide de l'API Vertex pour Python.

Revenez à votre instance de notebook et créez un notebook TensorFlow 2 via le lanceur :

new_notebook

Importez le SDK Vertex AI.

from google.cloud import aiplatform

Pour lancer le job d'entraînement à plusieurs nœuds de calcul, vous devez d'abord définir les spécifications du pool de nœuds de calcul. Notez que l'utilisation des GPU indiqués dans les spécifications est totalement facultative. Vous pouvez supprimer accelerator_type et accelerator_count si vous souhaitez utiliser un cluster de processeur uniquement, comme indiqué dans la section précédente.

# The spec of the worker pools including machine type and Docker image
# Be sure to replace {YOUR-PROJECT-ID} with your project ID.
worker_pool_specs=[
     {
        "replica_count": 1,
        "machine_spec": {
          "machine_type": "n1-standard-8", "accelerator_type": "NVIDIA_TESLA_V100", "accelerator_count": 1
        },
        "container_spec": {"image_uri": "gcr.io/{YOUR-PROJECT-ID}/multiworker:cassava"}
      },
      {
        "replica_count": 1,
        "machine_spec": {
          "machine_type": "n1-standard-8", "accelerator_type": "NVIDIA_TESLA_V100", "accelerator_count": 1
        },
        "container_spec": {"image_uri": "gcr.io/{YOUR-PROJECT-ID}/multiworker:cassava"}
      }
]

Ensuite, créez et exécutez un CustomJob. Vous devez remplacer {YOUR_BUCKET} par un bucket de votre projet pour la préproduction. Vous pouvez utiliser le bucket que vous avez créé précédemment.

# Replace YOUR_BUCKET
my_multiworker_job = aiplatform.CustomJob(display_name='multiworker-cassava-sdk',
                              worker_pool_specs=worker_pool_specs,
                              staging_bucket='gs://{YOUR_BUCKET}')

my_multiworker_job.run()

Votre job d'entraînement s'affiche dans l'onglet JOBS PERSONNALISÉS de la section "Entraînement" de votre console:

Tâches personnalisées

8. Nettoyage

Comme le notebook est configuré pour expirer au bout de 60 minutes d'inactivité, il n'est pas nécessaire d'arrêter l'instance. Si vous souhaitez arrêter l'instance manuellement, cliquez sur le bouton "Arrêter" dans la section "Vertex AI Workbench" de la console. Si vous souhaitez supprimer le notebook définitivement, cliquez sur le bouton "Supprimer".

Arrêter l'instance

Pour supprimer le bucket de stockage, utilisez le menu de navigation de la console Cloud pour accéder à Stockage, sélectionnez votre bucket puis cliquez sur "Supprimer" :

Supprimer l'espace de stockage