Guía de inicio rápido de PyTorch para VM de Cloud TPU

Descripción general: En esta guía de inicio rápido, se proporciona una breve introducción al trabajo con PyTorch y Cloud TPU.

Antes de comenzar

Antes de comenzar una de las guías de inicio rápido de una VM de Cloud TPU, asegúrate de leer la guía del usuario de Cloud TPU. Se describe cómo configurar un proyecto de Google Cloud y cómo crear una VM de Cloud TPU.

Accede a tu Cuenta de Google. Si aún no tienes una cuenta, regístrate para obtener una nueva. En Google Cloud Console, selecciona o crea un proyecto de Cloud desde la página de selección de proyecto. Asegúrate de tener habilitada la facturación para tu proyecto.

Habilita la API de Cloud TPU

Habilita la API de Cloud TPU con el siguiente comando de gcloud en Cloud Shell. (También puedes habilitarlo desde Google Cloud Console).

$ gcloud services enable tpu.googleapis.com

Configura el comando de gcloud

Ejecuta los comandos siguientes a fin de configurar gcloud a fin de usar tu proyecto de GCP y, luego, instalar los componentes necesarios para la vista previa de la VM de TPU.

$ gcloud config set account your-email-account
$ gcloud config set project project-id

Crea una VM de Cloud TPU con gcloud

$ gcloud alpha compute tpus tpu-vm create tpu-name \
  --zone=us-central1-b \
  --accelerator-type=v3-8 \
  --version=v2-alpha

Campos obligatorios

zone
La zona en la que planeas crear tu Cloud TPU.
accelerator-type
El tipo de Cloud TPU que se creará.
version
La versión del entorno de ejecución de Cloud TPU.

Conéctate a tu VM de Cloud TPU

$ gcloud alpha compute tpus tpu-vm ssh tpu-name --zone us-central1-b --project project-id

Campos obligatorios

tpu_name
El nombre de la VM de TPU a la que te conectas.
zone
La zona en la que creas el Cloud TPU. Actualmente, se admiten europe-west4-a y us-central1-a.
project-id
El ID de tu proyecto de GCP.

Establece la configuración del dispositivo de TPU XRT

  (vm)$ export XRT_TPU_CONFIG="localservice;0;localhost:51011"

Cambia la versión de PyTorch

Si prefiere usar PyTorch 1.9 en lugar de la versión preinstalada, ejecute estos comandos antes de comenzar el entrenamiento:

  (vm)$ sudo bash /var/scripts/docker-login.sh
  (vm)$ sudo docker rm libtpu || true
  (vm)$ sudo docker create --name libtpu gcr.io/cloud-tpu-v2-images/libtpu:pytorch-1.9 "/bin/bash"
  (vm)$ sudo docker cp libtpu:libtpu.so /lib
  (vm)$ sudo pip3 uninstall --yes torch torch_xla torchvision
  (vm)$ sudo pip3 install torch==1.9.0
  (vm)$ sudo pip3 install torchvision==0.10.0
  (vm)$ sudo pip3 install https://storage.googleapis.com/tpu-pytorch/wheels/tpuvm/torch_xla-1.9-cp38-cp38-linux_x86_64.whl

Realiza un cálculo simple:

Inicia el intérprete de Python 3:

(vm)$ python3

Resultado del comando:

Python 3.6.9 (default, Oct  8 2020, 12:12:24)
[GCC 8.4.0] on linux
Type "help", "copyright", "credits" or "license" for more information.

Importa paquetes de PyTorch

>>> import torch
>>> import torch_xla.core.xla_model as xm

Realiza el cálculo

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

Resultado del comando:

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

Sal del intérprete de Python:

>>> exit()

Limpia

Sigue estos pasos para evitar que se apliquen cargos a tu cuenta de Google Cloud por los recursos que usaste en esta guía de inicio rápido.

  1. Desconéctate de la instancia de Compute Engine, si aún no lo hiciste:

    (vm)$ exit

    El mensaje ahora debería mostrar username@projectname, que indica que estás en Cloud Shell.

  2. Borra tu Cloud TPU.

    $ gcloud alpha compute tpus tpu-vm delete tpu-name \
      --zone=us-central1-b
  3. Verifica que se hayan borrado los recursos con la ejecución de gcloud alpha compute tpus tpu-vm list. La eliminación puede tomar varios minutos. Una respuesta como la que se muestra a continuación indica que tus instancias se borraron de forma correcta.

     $ gcloud alpha compute tpus tpu-vm list --zone=us-central1-b

     NAME             STATUS

¿Qué sigue?

Obtén más información sobre las VM de Cloud TPU: