Entraînement Cloud TPU v5e

Cloud TPU v5e est l'accélérateur d'IA de dernière génération de Google Cloud. Avec un d'encombrement de 256 puces par pod, la version v5e est optimisée pour offrir produit pour les modèles Transformer, Text-to-Image et de réseaux de neurones convolutifs (CNN) l'entraînement, l'affinage et l'inférence. Pour en savoir plus sur l'utilisation de Cloud TPU Pour l'inférence, consultez la page Inférence avec v5e.

Pour en savoir plus sur le matériel et les configurations des TPU Cloud TPU v5e, consultez la page TPU v5e.

Commencer

Les sections suivantes décrivent comment commencer à utiliser TPU v5e.

Quota de requêtes

Vous avez besoin d'un quota d'utilisation de TPU v5e pour l'entraînement. Il existe différents types de quotas les TPU à la demande, les TPU réservés et les VM TPU Spot. Il existe des les quotas requis si vous utilisez votre TPU v5e pour inférence. Pour en savoir plus sur les quotas, consultez Quotas. Pour demander un quota de TPU v5e, contactez Cloud Ventes.

Créer un projet et un compte Google Cloud

Vous devez disposer d'un compte et d'un projet Google Cloud pour utiliser Cloud TPU. Pour plus consultez la page Configurer une instance Cloud TPU environnement.

Créer une instance Cloud TPU

Il est recommandé de provisionner les Cloud TPU v5es en file d'attente des ressources à l'aide de l'queued-resource create . Pour en savoir plus, consultez la section Gérer les fichiers ressources.

Vous pouvez également utiliser l'API Create Node (gcloud compute tpus tpu-vm create). pour provisionner Cloud TPU v5. Pour en savoir plus, consultez la page Gérer les ressources ressources.

Pour en savoir plus sur les métriques pour les configurations v5e pour l'entraînement, consultez la section Types de Cloud TPU v5e pour entraînement.

Configuration du framework

Cette section décrit le processus général de configuration de l'entraînement de modèle personnalisé à l'aide de JAX ou PyTorch avec TPU v5e. TensorFlow est compatible avec TPU tpu-vm-tf-2.17.0-pjrt et tpu-vm-tf-2.17.0-pod-pjrt versions d'exécution.

Pour obtenir des instructions sur la configuration de l'inférence, consultez la page Inférence v5e présentation.

Configuration pour JAX

Si vous avez des formes de tranches supérieures à huit chips, vous aurez plusieurs VM en une. tranche. Dans ce cas, vous devez utiliser l'option --worker=all pour exécuter l'installation sur toutes les VM TPU en une seule étape, sans utiliser SSH pour se connecter à chacune séparément:

gcloud compute tpus tpu-vm ssh ${TPU_NAME}  \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='pip install -U "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'

Description des options de commande

Variable Description
TPU_NAME ID textuel attribué par l'utilisateur au TPU, créé lors de la la demande de ressource en file d'attente est allouée.
PROJECT_ID Nom du projet Google Cloud. Utilisez un projet existant ou créez-en un. à <ph type="x-smartling-placeholder"></ph> Configurer votre projet Google Cloud
ZONE Voir les régions et zones TPU pour les zones compatibles.
Worker [class name, see definition and ref site provided] VM TPU ayant accès aux TPU sous-jacents.

Vous pouvez exécuter la commande suivante pour vérifier le nombre d'appareils (les résultats ont été produits ici, avec un segment "v5litepod-16"). Ce code vérifie que tout est installé correctement en vérifiant que JAX voit les TensorCores de Cloud TPU et peut exécuter des opérations de base:

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='python3 -c "import jax; print(jax.device_count()); print(jax.local_device_count())"'

Le résultat doit ressembler à ce qui suit :

SSH: Attempting to connect to worker 0...
SSH: Attempting to connect to worker 1...
SSH: Attempting to connect to worker 2...
SSH: Attempting to connect to worker 3...
16
4
16
4
16
4
16
4

jax.device_count() indique le nombre total de chips dans la tranche donnée. jax.local_device_count() indique le nombre de chips accessibles par un seul VM de cette tranche.

# Check the number of chips in the given slice by summing the count of chips
# from all VMs through the
# jax.local_device_count() API call.
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='python3 -c "import jax; xs=jax.numpy.ones(jax.local_device_count()); print(jax.pmap(lambda x: jax.lax.psum(x, \"i\"), axis_name=\"i\")(xs))"'

Le résultat doit ressembler à ce qui suit :

SSH: Attempting to connect to worker 0...
SSH: Attempting to connect to worker 1...
SSH: Attempting to connect to worker 2...
SSH: Attempting to connect to worker 3...
[16. 16. 16. 16.]
[16. 16. 16. 16.]
[16. 16. 16. 16.]
[16. 16. 16. 16.]

Essayez les tutoriels JAX de ce document pour commencer à Entraînement v5e à l'aide de JAX.

Configuration pour PyTorch

Notez que la version v5e n'est compatible qu'avec l'environnement d'exécution PJRT. et PyTorch 2.1+ utilisera PJRT comme environnement d'exécution par défaut pour toutes les versions de TPU.

Cette section explique comment commencer à utiliser PJRT sur v5e avec PyTorch/XLA avec pour tous les nœuds de calcul.

Installer des dépendances

gcloud compute tpus tpu-vm ssh ${TPU_NAME}  \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
      sudo apt-get update -y
      sudo apt-get install libomp5 -y
      pip3 install mkl mkl-include
      pip3 install tf-nightly tb-nightly tbp-nightly
      pip3 install numpy
      sudo apt-get install libopenblas-dev -y
      pip3 install torch~=2.1.0 torchvision torch_xla[tpu]~=2.1.0 -f https://storage.googleapis.com/libtpu-releases/index.html
      pip3 install torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html'

Si vous recevez un message d'erreur lors de l'installation des roues pour torch, torch_xla ou torchvision "J'aime" pkg_resources.extern.packaging.requirements.InvalidRequirement: Expected end or semicolon (after name and no valid version specifier) torch==nightly+20230222, rétrograder votre version à l'aide de la commande suivante:

pip3 install setuptools==62.1.0

Exécuter un script avec PJRT

unset LD_PRELOAD

Voici un exemple utilisant un script Python pour effectuer un calcul sur une VM v5e:

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker all \
   --command='
      export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$HOME/.local/lib/
      export PJRT_DEVICE=TPU_C_API
      export PT_XLA_DEBUG=0
      export USE_TORCH=ON
      unset LD_PRELOAD
      export TPU_LIBRARY_PATH=$HOME/.local/lib/python3.10/site-packages/libtpu/libtpu.so
      python3 -c "import torch; import torch_xla; import torch_xla.core.xla_model as xm; print(xm.xla_device()); dev = xm.xla_device(); t1 = torch.randn(3,3,device=dev); t2 = torch.randn(3,3,device=dev); print(t1 + t2)"'

Un résultat semblable à celui-ci doit s'afficher :

SSH: Attempting to connect to worker 0...
SSH: Attempting to connect to worker 1...
xla:0
tensor([[ 1.8611, -0.3114, -2.4208],
[-1.0731, 0.3422, 3.1445],
[ 0.5743, 0.2379, 1.1105]], device='xla:0')
xla:0
tensor([[ 1.8611, -0.3114, -2.4208],
[-1.0731, 0.3422, 3.1445],
[ 0.5743, 0.2379, 1.1105]], device='xla:0')

Essayez les tutoriels PyTorch de ce document pour commencer à Entraînement v5e à l'aide de PyTorch.

Supprimez votre TPU et la ressource en file d'attente à la fin de votre session. Pour supprimer un ressource en file d'attente, supprimez la tranche, puis la ressource en file d'attente en deux étapes:

gcloud compute tpus tpu-vm delete ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

Ces deux étapes peuvent également être utilisées pour supprimer les demandes de ressources en file d’attente qui se trouvent dans l'état FAILED.

Exemples JAX/FLAX

Les sections suivantes décrivent des exemples d'entraînement des modèles JAX et FLAX sur TPU v5e.

Entraîner ImageNet sur v5e

Ce tutoriel explique comment entraîner ImageNet sur v5e à l'aide de fausses données d'entrée. Si Si vous souhaitez utiliser des données réelles, consultez le fichier README sur GitHub.

Prérequis

  1. Créez des variables d'environnement :

    export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5litepod-16
export ZONE=us-west4-a
export RUNTIME_VERSION=v2-alpha-tpuv5-lite
export SERVICE_ACCOUNT=your_service_account
export TPU_NAME=your_tpu_name
export QUEUED_RESOURCE_ID=your_queued_resource_id
export QUOTA_TYPE=quota_type
export VALID_UNTIL_DURATION=1d

  2. Créez une ressource TPU:

    gcloud compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
   --node-id=${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --accelerator-type=${ACCELERATOR_TYPE} \
   --runtime-version=${RUNTIME_VERSION} \
   --valid-until-duration=${VALID_UNTIL_DURATION} \
   --service-account=${SERVICE_ACCOUNT} \
   --${QUOTA_TYPE}

    Vous pourrez vous connecter en SSH à votre VM TPU une fois que la ressource en file d'attente se trouvera l'état ACTIVE:

    gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}  \
   --project=${PROJECT_ID} \
   --zone=${ZONE}

    Lorsque l'état de QueuedResource est ACTIVE, la sortie doit ressembler à ce qui suit:

     state: ACTIVE

  3. Installez la dernière version de JAX et de jaxlib:

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='pip install -U "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'

  4. Clonez le modèle ImageNet et installez la configuration requise correspondante:

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='git clone https://github.com/google/flax.git && cd flax/examples/imagenet && pip install -r requirements.txt && pip install flax==0.7.4'

  5. Pour générer de fausses données, le modèle a besoin d'informations sur les dimensions le jeu de données. Ces informations peuvent être collectées à partir des métadonnées de l'ensemble de données ImageNet:

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='mkdir -p $HOME/flax/.tfds/metadata/imagenet2012/5.1.0 && curl https://raw.githubusercontent.com/tensorflow/datasets/v4.4.0/tensorflow_datasets/testing/metadata/imagenet2012/5.1.0/dataset_info.json --output $HOME/flax/.tfds/metadata/imagenet2012/5.1.0/dataset_info.json'

Entraîner le modèle

Une fois toutes les étapes précédentes terminées, vous pouvez entraîner le modèle.

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='cd flax/examples/imagenet && JAX_PLATFORMS=tpu python3 imagenet_fake_data_benchmark.py'

Supprimer le TPU et la ressource en file d'attente

Supprimez votre TPU et la ressource en file d'attente à la fin de votre session.

gcloud compute tpus tpu-vm delete ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

Modèles FLAX Hugging Face

Les modèles Hugging Face qui sont implémentés en FLAX à partir de sur Cloud TPU v5e. Cette section fournit des instructions pour exécuter modèles populaires.

Entraîner ViT sur Imagenette

Ce tutoriel vous montre comment entraîner le modèle Vision Transformer (ViT) de HuggingFace à l'aide de l'outil Fast AI Imagenette sur Cloud TPU v5e.

Le modèle ViT a été le premier à avoir entraîné avec succès un encodeur Transformer sur ImageNet, avec d'excellents résultats par rapport aux réseaux convolutifs. Pour plus consultez les ressources suivantes:

Prérequis

  1. Créez des variables d'environnement :

    export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5litepod-16
export ZONE=us-west4-a
export RUNTIME_VERSION=v2-alpha-tpuv5-lite
export SERVICE_ACCOUNT=your_service_account
export TPU_NAME=your_tpu_name
export QUEUED_RESOURCE_ID=your_queued_resource_id
export QUOTA_TYPE=quota_type
export VALID_UNTIL_DURATION=1d

  2. Créez une ressource TPU:

    gcloud compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
   --node-id=${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --accelerator-type=${ACCELERATOR_TYPE} \
   --runtime-version=${RUNTIME_VERSION} \
   --valid-until-duration=${VALID_UNTIL_DURATION} \
   --service-account=${SERVICE_ACCOUNT} \
   --${QUOTA_TYPE}

    Vous pourrez vous connecter en SSH à votre VM TPU une fois que votre ressource en file d'attente possède l'état ACTIVE:

    gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}  \
   --project=${PROJECT_ID} \
   --zone=${ZONE}

    Lorsque la ressource en file d'attente est à l'état ACTIVE, le résultat est semblable comme suit:

     state: ACTIVE

  3. Installez JAX et sa bibliothèque:

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='pip install -U "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'

  4. Télécharger le dépôt Hugging Face et les conditions d'installation requises:

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='git clone https://github.com/huggingface/transformers.git && cd transformers && pip install . && pip install -r examples/flax/_tests_requirements.txt && pip install --upgrade huggingface-hub urllib3 zipp && pip install tensorflow==2.17.0 && pip install -r examples/flax/vision/requirements.txt'

  5. Téléchargez l'ensemble de données Imagenette:

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='cd transformers && wget https://s3.amazonaws.com/fast-ai-imageclas/imagenette2.tgz && tar -xvzf imagenette2.tgz'

Entraîner le modèle

Entraînez le modèle avec un tampon prémappé de 4 Go.

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='cd transformers && JAX_PLATFORMS=tpu python3 examples/flax/vision/run_image_classification.py --train_dir "imagenette2/train" --validation_dir "imagenette2/val" --output_dir "./vit-imagenette" --learning_rate 1e-3 --preprocessing_num_workers 32 --per_device_train_batch_size 8 --per_device_eval_batch_size 8 --model_name_or_path google/vit-base-patch16-224-in21k --num_train_epochs 3'

Supprimer le TPU et la ressource en file d'attente

Supprimez votre TPU et la ressource en file d'attente à la fin de votre session.

gcloud compute tpus tpu-vm delete ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

Résultats de l'analyse comparative ViT

Le script d'entraînement a été exécuté sur v5litepod-4, v5litepod-16 et v5litepod-64. La Le tableau suivant indique les débits avec différents types d'accélérateurs.

Type d'accélérateur v5litepod-4 V5litepod-16 V5litepod-64
Époque 3 3 3
Taille du lot global 32 128 512
Débit (exemples/s) 263,40 429,34 470,71

Entraînement Diffusion sur Pokémon

Ce tutoriel vous explique comment entraîner le modèle Stable Diffusion à partir de HuggingFace utilisant le Pokémon sur Cloud TPU v5e.

Le modèle Stable Diffusion est un modèle texte-image latent qui génère des images réalistes à partir de n'importe quelle entrée de texte. Pour en savoir plus, consultez les ressources suivantes :

Prérequis

  1. Configurez un bucket de stockage pour la sortie du modèle. 

    gcloud storage buckets create gs://your_bucket 
    
--project=your_project 
    
--location=us-west1 
    
export GCS_BUCKET_NAME=your_bucket

  2. Créez des variables d'environnement

    export GCS_BUCKET_NAME=your_bucket
export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5litepod-16
export ZONE=us-west1-c
export RUNTIME_VERSION=v2-alpha-tpuv5-lite
export SERVICE_ACCOUNT=your_service_account
export TPU_NAME=your_tpu_name
export QUEUED_RESOURCE_ID=queued_resource_id
export QUOTA_TYPE=quota_type
export VALID_UNTIL_DURATION=1d

    Description des options de commande

    Variable Description
    GCS_BUCKET_NAME Affichage dans la console Google Cloud -> Cloud Storage -> Buckets
    PROJECT_ID Nom du projet Google Cloud. Utilisez un projet existant ou créez-en un. à <ph type="x-smartling-placeholder"></ph> Configurer votre projet Google Cloud
    ACCELERATOR_TYPE Consultez les versions TPU correspondant à votre version de TPU.
    ZONE Voir les régions et zones TPU pour les zones compatibles.
    RUNTIME_VERSION Utilisez v2-alpha-tpuv5 pour RUNTIME_VERSION.
    SERVICE_ACCOUNT Il s'agit de l'adresse de votre compte de service, dans la console Google Cloud -> IAM -> Comptes de service. Par exemple: [email protected]
    TPU_NAME ID textuel attribué par l'utilisateur au TPU, créé lors de la la demande de ressource en file d'attente est allouée.
    QUEUED_RESOURCE_ID ID du texte attribué par l'utilisateur dans la requête de ressource en file d'attente. Voir le document Queued Resources (Ressources en file d'attente) pour en savoir plus sur les ressources en file d'attente.
    QUOTA_TYPE Peut être reserved ou spot. Si aucune de ces options n'est spécifiée, le QUOTA_TYPE la valeur par défaut est on-demand. Voir les quotas pour en savoir plus sur les différents types de quotas compatibles avec Cloud TPU.
    VALID_UNTIL_DURATION Durée de validité de la requête. Voir <ph type="x-smartling-placeholder"></ph> Ressources en file d'attente pour connaître les différentes durées valides.

  3. Créez une ressource TPU:

    gcloud compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
   --node-id=${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --accelerator-type=${ACCELERATOR_TYPE} \
   --runtime-version=${RUNTIME_VERSION} \
   --valid-until-duration=${VALID_UNTIL_DURATION} \
   --service-account=${SERVICE_ACCOUNT} \
   --${QUOTA_TYPE}

    Vous pourrez vous connecter en SSH à votre VM TPU une fois que la ressource en file d'attente passe à l'état ACTIVE:

    gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}  \
   --project=${PROJECT_ID} \
   --zone=${ZONE}

    Lorsque la ressource en file d'attente est à l'état ACTIVE, le résultat ressemble à ce qui suit:

     state: ACTIVE

  4. Installez JAX et sa bibliothèque.

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='pip install "jax[tpu]==0.4.16" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'

  5. Télécharger le dépôt HuggingFace et les exigences d'installation.

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='git clone https://github.com/RissyRan/diffusers.git && cd diffusers && pip install . && pip install tensorflow==2.17.0 clu && pip install -U -r examples/text_to_image/requirements_flax.txt'

Entraîner le modèle

Entraînez le modèle avec un tampon prémappé de 4 Go.

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} --zone=${ZONE} --project ${PROJECT_ID} --worker=all --command="
   git clone https://github.com/google/maxdiffusion
   cd maxdiffusion
   git reset --hard 57629bcf4fa32fe5a57096b60b09f41f2fa5c35d # This identifies the GitHub commit to use.
   pip3 install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
   pip3 install -r requirements.txt
   pip3 install .
   export LIBTPU_INIT_ARGS=""
   python -m src.maxdiffusion.models.train src/maxdiffusion/configs/base_2_base.yml run_name=your_run base_output_directory=gs://${GCS_BUCKET_NAME}/ enable_profiler=False"

Supprimer le TPU et la ressource en file d'attente

Supprimez votre TPU et la ressource en file d'attente à la fin de votre session.

gcloud compute tpus tpu-vm delete ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

Résultats de l'analyse comparative de la diffusion

Le script d'entraînement a été exécuté sur les versions v5litepod-4, v5litepod-16 et v5litepod-64. Le tableau suivant présente les débits.

Type d'accélérateur v5litepod-4 V5litepod-16 V5litepod-64
Étape d'entraînement 1500 1500 1500
Taille du lot global 32 64 128
Débit (exemples/s) 36,53 43,71 49,36

Entraîner GPT2 sur l'ensemble de données OSCAR

Ce tutoriel explique comment entraîner le modèle GPT2 depuis HuggingFace en utilisant l'ensemble de données OSCAR sur Cloud TPU v5e.

GPT2 est un modèle Transformer pré-entraîné sur du texte brut, l'étiquetage manuel. Il a été entraîné pour prédire le mot suivant dans des phrases. Pour en savoir plus, consultez les ressources suivantes :

Prérequis

  1. Créez des variables d'environnement :

    export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5litepod-16
export ZONE=us-west4-a
export RUNTIME_VERSION=v2-alpha-tpuv5-lite
export SERVICE_ACCOUNT=your_service_account
export TPU_NAME=your_tpu_name
export QUEUED_RESOURCE_ID=queued_resource_id
export QUOTA_TYPE=quota_type
export VALID_UNTIL_DURATION=1d

  2. Créez une ressource TPU:

    gcloud compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
   --node-id=${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --accelerator-type=${ACCELERATOR_TYPE} \
   --runtime-version=${RUNTIME_VERSION} \
   --valid-until-duration=${VALID_UNTIL_DURATION} \
   --service-account=${SERVICE_ACCOUNT} \
   --${QUOTA_TYPE}

    Vous pourrez vous connecter en SSH à votre VM TPU une fois que votre ressource en file d'attente possède l'état ACTIVE:

    gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE}

    Lorsque la ressource en file d'attente est à l'état ACTIVE, le résultat est semblable comme suit:

     state: ACTIVE

  3. Installez JAX et sa bibliothèque.

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='pip install -U "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'

  4. Télécharger le dépôt HuggingFace et les exigences d'installation.

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='git clone https://github.com/huggingface/transformers.git && cd transformers && pip install . && pip install -r examples/flax/_tests_requirements.txt && pip install --upgrade huggingface-hub urllib3 zipp && pip install TensorFlow && pip install -r examples/flax/language-modeling/requirements.txt'

  5. Téléchargez les configurations pour entraîner le modèle.

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='cd transformers/examples/flax/language-modeling && gcloud storage cp gs://cloud-tpu-tpuvm-artifacts/v5litepod-preview/jax/gpt . --recursive'

Entraîner le modèle

Entraînez le modèle avec un tampon prémappé de 4 Go.

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='cd transformers/examples/flax/language-modeling && TPU_PREMAPPED_BUFFER_SIZE=4294967296 JAX_PLATFORMS=tpu python3 run_clm_flax.py --output_dir=./gpt --model_type=gpt2 --config_name=./gpt --tokenizer_name=./gpt --dataset_name=oscar --dataset_config_name=unshuffled_deduplicated_no --do_train --do_eval --block_size=512 --per_device_train_batch_size=4 --per_device_eval_batch_size=4 --learning_rate=5e-3 --warmup_steps=1000 --adam_beta1=0.9 --adam_beta2=0.98 --weight_decay=0.01 --overwrite_output_dir --num_train_epochs=3 --logging_steps=500 --eval_steps=2500'

Supprimer le TPU et la ressource en file d'attente

Supprimez votre TPU et la ressource en file d'attente à la fin de votre session.

gcloud compute tpus tpu-vm delete ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

Résultats de l'analyse comparative pour GPT2

Le script d'entraînement a été exécuté sur les versions v5litepod-4, v5litepod-16 et v5litepod-64. Le tableau suivant présente les débits.

v5litepod-4 V5litepod-16 V5litepod-64
Époque 3 3 3
Taille du lot global 64 64 64
Débit (exemples/s) 74,60 72,97 72,62

PyTorch/XLA

Les sections suivantes décrivent des exemples d'entraînement des modèles PyTorch/XLA sur TPU v5e.

Entraîner ResNet à l'aide de l'environnement d'exécution PJRT

PyTorch/XLA migre de XRT vers PjRt depuis PyTorch 2.0 et versions ultérieures. Voici les Mise à jour des instructions de configuration de v5e pour les charges de travail d'entraînement PyTorch/XLA.

Prérequis

  1. Créez des variables d'environnement :

    export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5litepod-16
export ZONE=us-west4-a
export RUNTIME_VERSION=v2-alpha-tpuv5-lite
export SERVICE_ACCOUNT=your_service_account
export TPU_NAME=tpu-name
export QUEUED_RESOURCE_ID=queued_resource_id
export QUOTA_TYPE=quota_type
export VALID_UNTIL_DURATION=1d

  2. Créez une ressource TPU:

    gcloud compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
   --node-id=${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --accelerator-type=${ACCELERATOR_TYPE} \
   --runtime-version=${RUNTIME_VERSION} \
   --valid-until-duration=${VALID_UNTIL_DURATION} \
   --service-account=${SERVICE_ACCOUNT} \
   --{QUOTA_TYPE}

    Vous pourrez vous connecter en SSH à votre VM TPU une fois que votre QueuedResource possède l'état ACTIVE:

    gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE}

    Lorsque la ressource en file d'attente est à l'état ACTIVE, le résultat ressemble à ce qui suit:

     state: ACTIVE

  3. Installer des dépendances spécifiques à Torch/XLA

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
      sudo apt-get update -y
      sudo apt-get install libomp5 -y
      pip3 install mkl mkl-include
      pip3 install tf-nightly tb-nightly tbp-nightly
      pip3 install numpy
      sudo apt-get install libopenblas-dev -y
      pip3 install torch~=2.1.0 torchvision torch_xla[tpu]~=2.1.0 -f https://storage.googleapis.com/libtpu-releases/index.html
      pip3 install torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html'
Entraîner le modèle ResNet
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
      date
      export PJRT_DEVICE=TPU_C_API
      export PT_XLA_DEBUG=0
      export USE_TORCH=ON
      export XLA_USE_BF16=1
      export LIBTPU_INIT_ARGS=--xla_jf_auto_cross_replica_sharding
      export LD_LIBRARY_PATH=/usr/local/lib:$LD_LIBRARY_PATH
      export TPU_LIBRARY_PATH=$HOME/.local/lib/python3.10/site-packages/libtpu/libtpu.so
      git clone https://github.com/pytorch/xla.git
      cd xla/
      git reset --hard caf5168785c081cd7eb60b49fe4fffeb894c39d9
      python3 test/test_train_mp_imagenet.py --model=resnet50  --fake_data --num_epochs=1 —num_workers=16  --log_steps=300 --batch_size=64 --profile'

Supprimer le TPU et la ressource en file d'attente

Supprimez votre TPU et la ressource en file d'attente à la fin de votre session.

gcloud compute tpus tpu-vm delete ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet
Résultat de l'analyse comparative

Le tableau suivant présente les débits de benchmark.

Type d'accélérateur Débit (exemples/seconde)
v5litepod-4 4240 ex/s
V5litepod-16 10 810 ex/s
V5litepod-64 46 154 ex/s

Entraîner GPT2 sur v5e

Ce tutoriel explique comment exécuter GPT2 dans la version 5e à l'aide du dépôt HuggingFace sur PyTorch/XLA à l'aide de l'ensemble de données wikitext.

Prérequis

  1. Créez des variables d'environnement :

    export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5litepod-16
export ZONE=us-west4-a
export RUNTIME_VERSION=v2-alpha-tpuv5-lite
export SERVICE_ACCOUNT=your_service_account
export TPU_NAME=your_tpu_name
export QUEUED_RESOURCE_ID=queued_resource_id
export QUOTA_TYPE=quota_type
export VALID_UNTIL_DURATION=1d

  2. Créez une ressource TPU:

    gcloud compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
   --node-id=${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --accelerator-type=${ACCELERATOR_TYPE} \
   --runtime-version=${RUNTIME_VERSION} \
   --valid-until-duration=${VALID_UNTIL_DURATION} \
   --service-account=${SERVICE_ACCOUNT} \
   --${QUOTA_TYPE}

    Vous pourrez vous connecter en SSH à votre VM TPU une fois que votre QueuedResource sera ACTIVE:

    gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE}

    Lorsque la ressource en file d'attente est à l'état ACTIVE, la sortie doit ressembler à ce qui suit:

    state: ACTIVE

  3. Installez les dépendances PyTorch/XLA.

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
      sudo apt-get -y update
      sudo apt install -y libopenblas-base
      pip3 install torchvision
      pip3 uninstall -y torch
      pip3 install torch~=2.1.0 torchvision torch_xla[tpu]~=2.1.0 -f https://storage.googleapis.com/libtpu-releases/index.html
      pip3 install torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html'

  4. Télécharger le dépôt HuggingFace et les exigences d'installation.

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
      git clone https://github.com/pytorch/xla.git
      pip install --upgrade accelerate
      git clone https://github.com/huggingface/transformers.git
      cd transformers
      git checkout ebdb185befaa821304d461ed6aa20a17e4dc3aa2
      pip install .
      git log -1
      pip install datasets evaluate scikit-learn
      '

  5. Télécharger les configurations du modèle pré-entraîné

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
      gcloud storage cp gs://cloud-tpu-tpuvm-artifacts/config/xl-ml-test/pytorch/gpt2/my_config_2.json transformers/examples/pytorch/language-modeling/ --recursive
      gcloud storage cp gs://cloud-tpu-tpuvm-artifacts/config/xl-ml-test/pytorch/gpt2/fsdp_config.json transformers/examples/pytorch/language-modeling/'

Entraîner le modèle

Entraînez le modèle 2B en utilisant une taille de lot de 16.

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
      export PJRT_DEVICE=TPU_C_API
      cd transformers/
      export LD_LIBRARY_PATH=/usr/local/lib/
      export PT_XLA_DEBUG=0
      export USE_TORCH=ON
      python3 examples/pytorch/xla_spawn.py \
         --num_cores=4 \
         examples/pytorch/language-modeling/run_clm.py \
         --num_train_epochs=3 \
         --dataset_name=wikitext \
         --dataset_config_name=wikitext-2-raw-v1 \
         --per_device_train_batch_size=16 \
         --per_device_eval_batch_size=16 \
         --do_train \
         --do_eval \
         --logging_dir=./tensorboard-metrics \
         --cache_dir=./cache_dir \
         --output_dir=/tmp/test-clm \
         --overwrite_output_dir \
         --cache_dir=/tmp \
         --config_name=examples/pytorch/language-modeling/my_config_2.json \
         --tokenizer_name=gpt2 \
         --block_size=1024 \
         --optim=adafactor \
         --adafactor=true \
         --save_strategy=no \
         --logging_strategy=no \
         --fsdp=full_shard \
         --fsdp_config=examples/pytorch/language-modeling/fsdp_config.json'

Supprimer le TPU et la ressource en file d'attente

Supprimez votre TPU et la ressource en file d'attente à la fin de votre session.

gcloud compute tpus tpu-vm delete ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

Résultat de l'analyse comparative

Le script d'entraînement a été exécuté sur les versions v5litepod-4, v5litepod-16 et v5litepod-64. La Le tableau suivant montre les débits de benchmark pour les différents types d'accélérateurs.

v5litepod-4 V5litepod-16 V5litepod-64
Époque 3 3 3
config 600M 2 milliards 16 Mrds
Taille du lot global 64 128 256
Débit (exemples/s) 66 77 31

Entraîner ViT sur v5e

Ce tutoriel explique comment exécuter VIT sur la version 5e à l'aide du dépôt HuggingFace. sur PyTorch/XLA sur l'ensemble de données cifar10.

Prérequis

  1. Créez des variables d'environnement :

    export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5litepod-16
export ZONE=us-west4-a
export RUNTIME_VERSION=v2-alpha-tpuv5-lite
export SERVICE_ACCOUNT=your_service_account
export TPU_NAME=tpu-name
export QUEUED_RESOURCE_ID=queued_resource_id
export QUOTA_TYPE=quota_type
export VALID_UNTIL_DURATION=1d

  2. Créez une ressource TPU:

    gcloud compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
   --node-id=${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --accelerator-type=${ACCELERATOR_TYPE} \
   --runtime-version=${RUNTIME_VERSION} \
   --valid-until-duration=${VALID_UNTIL_DURATION} \
   --service-account=${SERVICE_ACCOUNT} \
   --${QUOTA_TYPE}

    Vous pourrez vous connecter en SSH à votre VM TPU une fois que votre QueuedResource sera l'état ACTIVE:

     gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE}

    Lorsque la ressource en file d'attente est à l'état ACTIVE, la sortie doit ressembler à ce qui suit:

     state: ACTIVE

  3. Installer des dépendances PyTorch/XLA

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all
   --command='
      sudo apt-get update -y
      sudo apt-get install libomp5 -y
      pip3 install mkl mkl-include
      pip3 install tf-nightly tb-nightly tbp-nightly
      pip3 install numpy
      sudo apt-get install libopenblas-dev -y
      pip3 install torch~=2.1.0 torchvision torch_xla[tpu]~=2.1.0 -f https://storage.googleapis.com/libtpu-releases/index.html
      pip3 install torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html'

  4. Télécharger le dépôt HuggingFace et les exigences d'installation.

       gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command="
      git clone https://github.com/suexu1025/transformers.git vittransformers; \
      cd vittransformers; \
      pip3 install .; \
      pip3 install datasets; \
      wget https://github.com/pytorch/xla/blob/master/scripts/capture_profile.py"

Entraîner le modèle

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
      export PJRT_DEVICE=TPU_C_API
      export PT_XLA_DEBUG=0
      export USE_TORCH=ON
      export TF_CPP_MIN_LOG_LEVEL=0
      export XLA_USE_BF16=1
      export LD_LIBRARY_PATH=/usr/local/lib:$LD_LIBRARY_PATH
      export TPU_LIBRARY_PATH=$HOME/.local/lib/python3.10/site-packages/libtpu/libtpu.so
      cd vittransformers
      python3 -u examples/pytorch/xla_spawn.py --num_cores 4 examples/pytorch/image-pretraining/run_mae.py --dataset_name=cifar10 \
      --remove_unused_columns=False \
      --label_names=pixel_values \
      --mask_ratio=0.75 \
      --norm_pix_loss=True \
      --do_train=true \
      --do_eval=true \
      --base_learning_rate=1.5e-4 \
      --lr_scheduler_type=cosine \
      --weight_decay=0.05 \
      --num_train_epochs=3 \
      --warmup_ratio=0.05 \
      --per_device_train_batch_size=8 \
      --per_device_eval_batch_size=8 \
      --logging_strategy=steps \
      --logging_steps=30 \
      --evaluation_strategy=epoch \
      --save_strategy=epoch \
      --load_best_model_at_end=True \
      --save_total_limit=3 \
      --seed=1337 \
      --output_dir=MAE \
      --overwrite_output_dir=true \
      --logging_dir=./tensorboard-metrics \
      --tpu_metrics_debug=true'

Supprimer le TPU et la ressource en file d'attente

Supprimez votre TPU et la ressource en file d'attente à la fin de votre session.

gcloud compute tpus tpu-vm delete ${TPU_NAME}
   --project=${PROJECT_ID}
   --zone=${ZONE}
   --quiet

gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID}
   --project=${PROJECT_ID}
   --zone=${ZONE}
   --quiet

Résultat de l'analyse comparative

Le tableau suivant indique les débits de benchmark pour les différents types d'accélérateurs.

v5litepod-4 V5litepod-16 V5litepod-64
Époque 3 3 3
Taille du lot global 32 128 512
Débit (exemples/s) 201 657 2 844

TensorFlow 2.x

Les sections suivantes décrivent des exemples d'entraînement de TensorFlow 2.x sur les TPU v5e.

Entraîner Resnet sur une version v5e à hôte unique

Ce tutoriel explique comment entraîner ImageNet sur v5litepod-4 ou v5litepod-8 à l'aide d'un faux ensemble de données. Si vous souhaitez utiliser un autre jeu de données, reportez-vous à Préparer l'ensemble de données

Prérequis

  1. Créez des variables d'environnement :

    export PROJECT_ID=your-project-ID
export ACCELERATOR_TYPE=v5litepod-4
export ZONE=us-east1-c
export RUNTIME_VERSION=tpu-vm-tf-2.15.0-pjrt
export TPU_NAME=your-tpu-name
export QUEUED_RESOURCE_ID=your-queued-resource-id
export QUOTA_TYPE=quota-type

    ACCELERATOR_TYPE peut être v5litepod-4 ou v5litepod-8.

  2. Créez une ressource TPU:

    gcloud compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
   --node-id=${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --accelerator-type=${ACCELERATOR_TYPE} \
   --runtime-version=${RUNTIME_VERSION} \
   --${QUOTA_TYPE}

    Vous pourrez vous connecter en SSH à votre VM TPU une fois que la ressource en file d'attente se trouvera l'état ACTIVE. Pour vérifier l'état de votre ressource en file d'attente, exécutez la commande suivante:

    gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE}

  3. Se connecter à votre TPU à l'aide de SSH

    gcloud compute tpus tpu-vm ssh ${TPU_NAME}  \
   --project=${PROJECT_ID} \
   --zone=${ZONE}

  4. Définir des variables d'environnement

    export MODELS_REPO=/usr/share/tpu/models
export PYTHONPATH="${MODELS_REPO}:${PYTHONPATH}"
export MODEL_DIR=gcp-directory-to-store-model
export DATA_DIR=gs://cloud-tpu-test-datasets/fake_imagenet
export NEXT_PLUGGABLE_DEVICE_USE_C_API=true
export TF_PLUGGABLE_DEVICE_LIBRARY_PATH=/lib/libtpu.so

  5. Accédez au répertoire du dépôt de modèles et installez les conditions requises.

    cd ${MODELS_REPO} && git checkout r2.15.0
pip install -r official/requirements.txt

Entraîner le modèle

Exécutez le script d'entraînement.

python3 official/vision/train.py \
   --tpu=local \
   --experiment=resnet_imagenet \
   --mode=train_and_eval \
   --config_file=official/vision/configs/experiments/image_classification/imagenet_resnet50_tpu.yaml \
   --model_dir=${MODEL_DIR} \
   --params_override="runtime.distribution_strategy=tpu,task.train_data.input_path=${DATA_DIR}/train*,task.validation_data.input_path=${DATA_DIR}/validation*,task.train_data.global_batch_size=2048,task.validation_data.global_batch_size=2048,trainer.train_steps=100"

Supprimer le TPU et la ressource en file d'attente

  1. Supprimer votre TPU

    gcloud compute tpus tpu-vm delete ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

  2. Supprimer votre requête de ressource en file d'attente

    gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

Entraîner Resnet sur une v5e à hôtes multiples

Ce tutoriel explique comment entraîner ImageNet sur v5litepod-16 ou version ultérieure à l'aide de un jeu de données factice. Si vous souhaitez utiliser un autre ensemble de données, consultez la section Préparer l'ensemble de données.

  1. Créez des variables d'environnement :

    export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5litepod-16
export ZONE=us-east1-c
export RUNTIME_VERSION=tpu-vm-tf-2.15.0-pod-pjrt
export TPU_NAME=your_tpu_name
export QUEUED_RESOURCE_ID=your-queued-resource-id
export QUOTA_TYPE=quota-type

    ACCELERATOR_TYPE peut être de taille v5litepod-16 ou supérieure.

  2. Créez une ressource TPU:

    gcloud compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
   --node-id=${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --accelerator-type=${ACCELERATOR_TYPE} \
   --runtime-version=${RUNTIME_VERSION} \
   --${QUOTA_TYPE}

    Vous pourrez vous connecter en SSH à votre VM TPU une fois que la ressource en file d'attente se trouvera l'état ACTIVE. Pour vérifier l'état d'une ressource en file d'attente, utilisez la méthode la commande suivante:

    gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE}

  3. Se connecter à votre TPU (nœud de calcul zéro) à l'aide de SSH

    gcloud compute tpus tpu-vm ssh ${TPU_NAME}  \
   --project=${PROJECT_ID} \
   --zone=${ZONE}

  4. Définir des variables d'environnement

    export MODELS_REPO=/usr/share/tpu/models
export PYTHONPATH="${MODELS_REPO}:${PYTHONPATH}"
export MODEL_DIR=gcp-directory-to-store-model
export DATA_DIR=gs://cloud-tpu-test-datasets/fake_imagenet
export TPU_LOAD_LIBRARY=0
export TPU_NAME=your_tpu_name

  5. Accédez au répertoire du dépôt de modèles et installez les conditions requises.

     cd $MODELS_REPO && git checkout r2.15.0
 pip install -r official/requirements.txt

Entraîner le modèle

Exécutez le script d'entraînement.

python3 official/vision/train.py \
   --tpu=${TPU_NAME} \
   --experiment=resnet_imagenet \
   --mode=train_and_eval \
   --model_dir=${MODEL_DIR} \
   --params_override="runtime.distribution_strategy=tpu,task.train_data.input_path=${DATA_DIR}/train*, task.validation_data.input_path=${DATA_DIR}/validation*"

Supprimer le TPU et la ressource en file d'attente

  1. Supprimer votre TPU

    gcloud compute tpus tpu-vm delete ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

  2. Supprimer votre requête de ressource en file d'attente

    gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet