Entraîner ResNet-50 sur Cloud TPU avec PyTorch

Attention : Cette page s'applique à l'API Cloud TPU, qui n'est plus en cours de développement. L'API Cloud TPU ne recevra que des correctifs de bugs et des mises à jour de sécurité. Pour bénéficier des dernières fonctionnalités et de la compatibilité avec les dernières versions de TPU, nous vous recommandons d'utiliser Compute Engine ou Google Kubernetes Engine pour gérer les ressources TPU. Pour en savoir plus, consultez Ressources TPU dans Compute Engine et À propos des TPU dans GKE. L'API Cloud TPU n'est pas compatible avec les TPU7x ni les versions ultérieures de TPU.

Ce tutoriel explique comment entraîner le modèle ResNet-50 sur un appareil Cloud TPU avec PyTorch. La même procédure peut s'appliquer à d'autres modèles de classification d'image optimisés pour TPU, qui utilisent PyTorch et l'ensemble de données ImageNet.

Le modèle utilisé dans ce tutoriel est basé sur l'article Deep Residual Learning for Image Recognition (Deep learning résiduel pour la reconnaissance d'images), qui présente l'architecture de réseau résiduel (ResNet). Le tutoriel emploie la variante à 50 couches, ResNet-50, et illustre l'entraînement du modèle à l'aide de PyTorch/XLA.

Attention : Ce tutoriel utilise un ensemble de données tiers. Google ne formule aucune déclaration ni garantie quant à la validité ou à tout autre aspect de cet ensemble de données.

Objectifs

• Préparer l'ensemble de données
• Exécuter la tâche d'entraînement
• Vérifier les résultats

Coûts

Dans ce document, vous utilisez les composants facturables suivants de Google Cloud :

• Compute Engine
• Cloud TPU

Pour obtenir une estimation des coûts en fonction de votre utilisation prévue, utilisez le simulateur de coût.

Les nouveaux utilisateurs de Google Cloud peuvent bénéficier d'un essai sans frais.

Avant de commencer

Avant de commencer ce tutoriel, vérifiez que votre projet Google Cloud est correctement configuré.

Connectez-vous à votre compte Google Cloud . Si vous débutez sur Google Cloud, créez un compte pour évaluer les performances de nos produits en conditions réelles. Les nouveaux clients bénéficient également de 300 $ de crédits sans frais pour exécuter, tester et déployer des charges de travail.

In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

Roles required to select or create a project

• Select a project: Selecting a project doesn't require a specific IAM role—you can select any project that you've been granted a role on.
• Create a project: To create a project, you need the Project Creator role (roles/resourcemanager.projectCreator), which contains the resourcemanager.projects.create permission. Learn how to grant roles.

Note: If you don't plan to keep the resources that you create in this procedure, create a project instead of selecting an existing project. After you finish these steps, you can delete the project, removing all resources associated with the project.

Go to project selector

Verify that billing is enabled for your Google Cloud project.

In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

Roles required to select or create a project

• Select a project: Selecting a project doesn't require a specific IAM role—you can select any project that you've been granted a role on.
• Create a project: To create a project, you need the Project Creator role (roles/resourcemanager.projectCreator), which contains the resourcemanager.projects.create permission. Learn how to grant roles.

Note: If you don't plan to keep the resources that you create in this procedure, create a project instead of selecting an existing project. After you finish these steps, you can delete the project, removing all resources associated with the project.

Go to project selector

Verify that billing is enabled for your Google Cloud project.

Ce tutoriel utilise des composants facturables de Google Cloud. Consultez la grille tarifaire de Cloud TPU pour estimer vos coûts. Veillez à nettoyer les ressources que vous avez créées lorsque vous avez terminé, afin d'éviter des frais inutiles.

Créez une VM TPU.

1. Ouvrez une fenêtre Cloud Shell ou de terminal.

   Ouvrir Cloud Shell

2. Définissez des variables d'environnement pour les paramètres de la commande Google Cloud CLI :

   export PROJECT_ID=your-project-id
   export TPU_NAME=your-tpu-name
   export ZONE=us-central1-b
   export ACCELERATOR_TYPE=v6e-8
   export RUNTIME_VERSION=v2-alpha-tpuv6e

   Descriptions des variables d'environnement

   • PROJECT_ID : ID de votre projet Google Cloud . Utilisez un projet existant ou créez-en un.
   • TPU_NAME : nom du TPU.
   • ZONE : Zone dans laquelle créer la VM TPU. Pour en savoir plus sur les zones compatibles, consultez Régions et zones TPU.
   • ACCELERATOR_TYPE : le type d'accélérateur spécifie la version et la taille du Cloud TPU que vous souhaitez créer. Pour en savoir plus sur les types d'accélérateurs compatibles avec chaque version de TPU, consultez Versions de TPU.
   • RUNTIME_VERSION : Version logicielle du Cloud TPU.

3. Créez une VM TPU.

    gcloud compute tpus tpu-vm create $TPU_NAME \
    --accelerator-type=$ACCELERATOR_TYPE \
    --version=$RUNTIME_VERSION \
    --zone=$ZONE \
    --project=$PROJECT_ID
   

   Remarque : La première fois que vous exécutez une commande dans une nouvelle VM Cloud Shell, une page Authorize Cloud Shell s'affiche. Cliquez sur Authorize en bas de la page pour autoriser gcloud à effectuer des appels d'API Google Cloud avec vos identifiants.

4. Connectez-vous à votre VM TPU via SSH :

    gcloud compute tpus tpu-vm ssh  $TPU_NAME --zone=$ZONE
   

Installer PyTorch/XLA sur votre VM TPU

1. Créez un fichier nommé requirements.txt et ajoutez-y le contenu suivant :

   --find-links https://storage.googleapis.com/libtpu-releases/index.html
   --find-links https://storage.googleapis.com/libtpu-wheels/index.html
   torch~=2.6.0
   torch_xla[tpu]~=2.6.0
   torchvision
   ray[default]==2.40.0
   

2. Installez les dépendances :

   pip install -r requirements.txt
   

3. Clonez le dépôt GitHub PyTorch/XLA.

   git clone --depth=1 https://github.com/pytorch/xla.git
   

4. Exécutez le script d'entraînement avec des données factices.

   PJRT_DEVICE=TPU python3 xla/test/test_train_mp_imagenet.py \
   --fake_data \
   --batch_size=256 \
   --num_epochs=1
   

Effectuer un nettoyage

Pour éviter que les ressources utilisées dans ce tutoriel soient facturées sur votre compte Google Cloud, supprimez le projet contenant les ressources, ou conservez le projet et supprimez chaque ressource individuellement.

1. Déconnectez-vous de la VM TPU :

   exit
   

2. Supprimez votre VM TPU.

   gcloud compute tpus tpu-vm delete $TPU_NAME \
       --zone=$ZONE
   

Étapes suivantes

• Entraîner des modèles de diffusion avec PyTorch
• Résoudre les problèmes liés à PyTorch sur les TPU
• Documentation PyTorch/XLA