Entraîner Resnet50 sur Cloud TPU avec PyTorch

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.

Objectifs

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

Coûts

Ce tutoriel utilise des composants facturables de Google Cloud, dont :

  • Instance
  • Cloud TPU

Obtenez une estimation des coûts en fonction de votre utilisation prévue à l'aide du simulateur de coût. Les nouveaux utilisateurs de Google Cloud peuvent bénéficier d'un essai gratuit.

Avant de commencer

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

  1. Connectez-vous à votre compte Google.

    Si vous n'en possédez pas déjà un, vous devez en créer un.

  2. Dans Cloud Console, sur la page de sélection du projet, sélectionnez ou créez un projet Cloud.

    Accéder à la page de sélection du projet

  3. Vérifiez que la facturation est activée pour votre projet Google Cloud. Découvrez comment vérifier que la facturation est activée pour votre projet.

  4. 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.

Configurer une instance Compute Engine

  1. Ouvrez une fenêtre Cloud Shell.

    Ouvrir Cloud Shell

  2. Créez une variable pour l'ID de votre projet.

    export PROJECT_ID=project-id
    
  3. Configurez l'outil de ligne de commande gcloud pour utiliser le projet dans lequel vous souhaitez créer Cloud TPU.

    gcloud config set project ${PROJECT_ID}
    
  4. Définir la version PyTorch

    export VERSION=1.6
    
  5. Lancez la ressource Compute Engine nécessaire pour ce tutoriel.

    gcloud compute instances create resnet50-tutorial \
    --zone=us-central1-a \
    --machine-type=n1-highmem-96  \
    --image-family=torch-xla \
    --image-project=ml-images  \
    --boot-disk-size=300GB \
    --scopes=https://www.googleapis.com/auth/cloud-platform
    

    Si vous prévoyez d'entraîner Resnet50 sur des données réelles, choisissez le type de machine comportant le plus grand nombre de processeurs possible. Resnet50 est généralement fortement subordonné aux entrées. L'entraînement peut donc être assez lent, à moins de disposer d'un grand nombre de nœuds de calcul pour alimenter les données et de suffisamment de RAM pour gérer un grand nombre de threads de calcul. Pour obtenir de meilleurs résultats, sélectionnez le type de machine n1-highmem-96.

    Si vous envisagez de télécharger ImageNet, spécifiez une taille de disque d'au moins 300 Go. Si vous prévoyez de n'utiliser que des données factices, vous pouvez spécifier la taille de disque par défaut de 20 Go. Pour ce tutoriel, il est conseillé d'utiliser les deux ensembles de données.

Lancer une ressource Cloud TPU

  1. À partir de la machine virtuelle Compute Engine, lancez une ressource Cloud TPU à l'aide de la commande suivante :

    (vm) $ gcloud compute tpus create resnet50-tutorial \
    --zone=us-central1-a \
    --network=default \
    --version=pytorch-${VERSION?}  \
    --accelerator-type=v3-8
    
  2. Identifiez l'adresse IP de la ressource Cloud TPU.

    (vm) $ gcloud compute tpus list --zone=us-central1-a
    

Créer et configurer l'environnement PyTorch

  1. Connectez-vous à la nouvelle instance Compute Engine

    gcloud compute ssh resnet50-tutorial --zone=us-central1-a
    

À partir de là, le préfixe (vm)$ signifie que vous devez exécuter la commande sur l'instance de VM Compute Engine.

  1. Démarrez un environnement conda.

    (vm) $ export VERSION=1.6
    (vm) $ conda activate torch-xla-${VERSION?}
    
  2. Configurez des variables d'environnement pour la ressource Cloud TPU.

    (vm) $ export TPU_IP_ADDRESS=ip-address
    
    (vm) $ export XRT_TPU_CONFIG="tpu_worker;0;$TPU_IP_ADDRESS:8470"
    

Entraîner avec l'ensemble de données factices (fake_data)

Pour votre première exécution, nous vous recommandons d'utiliser l'ensemble de données factice plutôt que l'ensemble de données réel ImageNet, car fake_data est automatiquement installé dans votre VM et nécessite moins de ressources.

(vm) $ python /usr/share/torch-xla-${VERSION?}/pytorch/xla/test/test_train_mp_imagenet.py --fake_data --model=resnet50 --num_epochs=2 --batch_size=128 --log_steps=20

Entraîner avec l'ensemble de données réel

Si tout semble correct avec l'indicateur --fake_data, vous pouvez commencer l'entraînement sur des données réelles, telles que ImageNet.

En général, test_train_mp_imagenet.py utilise torchvision.datasets.ImageFolder. Vous pouvez donc utiliser n'importe quel ensemble de données correctement structuré. Consultez la documentation sur ImageFolder.

Voici quelques suggestions de modification de la ligne de commande pour l'utilisation de données réelles, en supposant que vous avez stocké l'ensemble de données dans ~/imagenet :

(vm) $ python /usr/share/torch-xla-${VERSION?}/pytorch/xla/test/test_train_mp_imagenet.py --datadir=~/imagenet --model=resnet50 --num_epochs=90 --num_workers=8 --batch_size=128 --log_steps=200

Demander une augmentation de quota de processeurs

Effectuez votre demande de ressources supplémentaires quelques jours à l'avance pour que le temps de traitement soit suffisant.

  1. Accédez à la page Quotas.

    Accéder à la page Quotas

  2. Dans le menu Service, sélectionnez API Cloud TPU.
  3. Sélectionnez la région ou les zones dans lesquelles vous souhaitez utiliser les processeurs.
  4. Dans le menu Métrique, sélectionnez Aucun, puis saisissez les processeurs dans le champ de recherche.
  5. Sélectionnez Processeurs.
  6. Dans la liste, sélectionnez API Compute Engine – Processeurs, puis cliquez sur Modifier les quotas en haut de la page.
  7. Saisissez le montant du quota que vous demandez et une description (obligatoire), puis cliquez sur Terminé. Une demande est envoyée à votre fournisseur de services pour approbation.

Nettoyer

Pour éviter que les ressources utilisées dans ce tutoriel soient facturées sur votre compte Google Cloud Platform :

  1. Déconnectez-vous de l'instance Compute Engine, si vous ne l'avez pas déjà fait :

    (vm)$ exit
    

    Votre invite devrait maintenant être user@projectname, indiquant que vous êtes dans Cloud Shell.

  2. Dans votre environnement Cloud Shell, exécutez ctpu delete avec les indicateurs --zone et --name que vous avez utilisés lors de la configuration de la VM Compute Engine et du Cloud TPU. Cela supprime à la fois votre VM et votre Cloud TPU.

    $ ctpu delete --name=resnet50-tutorial \
      --zone=europe-west4-a
    
  3. Exécutez ctpu status pour vérifier qu'il ne reste aucune instance allouée afin d'éviter des frais inutiles liés à l'utilisation des ressources TPU. La suppression peut prendre plusieurs minutes. Si vous n'avez plus d'instances allouées, une réponse de ce type s'affiche :

    $ ctpu status --zone=europe-west4-a
    
    2018/04/28 16:16:23 WARNING: Setting zone to "europe-west4-a"
    No instances currently exist.
            Compute Engine VM:     --
            Cloud TPU:             --
    

Étape suivante

Essayez les colabs PyTorch :