Effectuer un calcul sur une VM Cloud TPU à l'aide de PyTorch

Ce guide de démarrage rapide explique comment créer un Cloud TPU, installer PyTorch et exécuter un calcul simple sur Cloud TPU. Pour un tutoriel plus détaillé montrant comment entraîner un modèle sur une ressource Cloud TPU, consultez l'un des tutoriels PyTorch pour Cloud TPU.

Avant de commencer

Avant de suivre ce guide de démarrage rapide, vous devez créer un compte Google Cloud Platform, installer Google Cloud CLI et configurer la commande gcloud. Pour en savoir plus, consultez la page Configurer un compte et un projet Cloud TPU.

Créer un Cloud TPU avec `gcloud`

Pour créer une VM TPU dans le projet utilisateur par défaut, le réseau et la zone de calcul sont exécutés:

$ gcloud compute tpus tpu-vm create tpu-name \
   --zone=us-central2-b \
   --accelerator-type=v3-8 \
   --version=tpu-vm-pt-2.0

Description des options de commande

zone: Zone dans laquelle vous prévoyez de créer votre Cloud TPU.
accelerator-type: Type du Cloud TPU à créer.
version: Version du logiciel Cloud TPU.

Lors de la création de votre TPU, vous pouvez transmettre les indicateurs --network et --subnetwork supplémentaires si vous souhaitez spécifier le réseau et le sous-réseau par défaut. Si vous ne souhaitez pas utiliser le réseau par défaut, vous devez transmettre l'option --network. L'option --subnetwork est facultative et peut être utilisée pour spécifier un sous-réseau par défaut pour le réseau que vous utilisez (par défaut ou spécifié par l'utilisateur). Pour en savoir plus sur ces options, consultez la page de référence de l'API gcloud.

Se connecter à la VM Cloud TPU

   $ gcloud compute tpus tpu-vm ssh tpu-name 

   --zone=us-central2-b

Définir la configuration de l'environnement d'exécution TPU

Assurez-vous que l'environnement d'exécution PyTorch/XLA utilise le TPU.

   (vm) $ export PJRT_DEVICE=TPU

Effectuez un calcul simple :

Créez un fichier nommé tpu-test.py dans le répertoire actuel et collez-y le script suivant.

import torch
import torch_xla.core.xla_model as xm

dev = xm.xla_device()
t1 = torch.randn(3,3,device=dev)
t2 = torch.randn(3,3,device=dev)
print(t1 + t2)

Exécutez le script :

  (vm)$ python3 tpu-test.py

Le résultat du script montre le résultat du calcul :

tensor([[-0.2121,  1.5589, -0.6951],
        [-0.7886, -0.2022,  0.9242],
        [ 0.8555, -1.8698,  1.4333]], device='xla:1')

Effectuer un nettoyage

Pour éviter que les ressources utilisées sur cette page soient facturées sur votre compte Google Cloud, procédez comme suit :

Déconnectez-vous de l'instance Compute Engine, si vous ne l'avez pas déjà fait :
```
(vm)$ exit
```
Votre invite devrait maintenant être username@projectname, indiquant que vous êtes dans Cloud Shell.

Supprimez votre Cloud TPU.

$ gcloud compute tpus tpu-vm delete tpu-name \
  --zone=us-central2-b

Le résultat de cette commande doit confirmer que votre TPU a été supprimé.

Étapes suivantes

En savoir plus sur les VM Cloud TPU :