Esta página se ha traducido con Cloud Translation API.

Ejecutar un cálculo en una VM de TPU de Cloud con PyTorch

En este documento se ofrece una breve introducción sobre cómo trabajar con PyTorch y Cloud TPU.

Antes de empezar

Antes de ejecutar los comandos de este documento, debes crear una Google Cloud cuentagcloud, instalar la CLI de Google Cloud y configurar el comando gcloud. Para obtener más información, consulta el artículo Configurar el entorno de TPU de Cloud.

Crear una TPU de Cloud con `gcloud`

Define algunas variables de entorno para que los comandos sean más fáciles de usar.

export PROJECT_ID=your-project-id
export TPU_NAME=your-tpu-name
export ZONE=us-east5-a
export ACCELERATOR_TYPE=v5litepod-8
export RUNTIME_VERSION=v2-alpha-tpuv5-lite

Descripciones de las variables de entorno

Variable	Descripción
`PROJECT_ID`	El ID de tu proyecto Google Cloud . Usa un proyecto que ya tengas o crea uno.
`TPU_NAME`	El nombre de la TPU.
`ZONE`	La zona en la que se creará la VM de TPU. Para obtener más información sobre las zonas admitidas, consulta Regiones y zonas de TPU.
`ACCELERATOR_TYPE`	El tipo de acelerador especifica la versión y el tamaño de la TPU de Cloud que quieres crear. Para obtener más información sobre los tipos de aceleradores compatibles con cada versión de TPU, consulta Versiones de TPU.
`RUNTIME_VERSION`	La versión de software de la TPU de Cloud.

Crea tu VM de TPU ejecutando el siguiente comando:

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

Conéctate a tu VM de TPU de Cloud

Conéctate a tu VM de TPU mediante SSH con el siguiente comando:

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

Si no puedes conectarte a una VM de TPU mediante SSH, puede deberse a que la VM de TPU no tiene una dirección IP externa. Para acceder a una VM de TPU sin una dirección IP externa, sigue las instrucciones que se indican en Conectarse a una VM de TPU sin una dirección IP pública.

Instalar PyTorch/XLA en tu VM de TPU

$ (vm) sudo apt-get update
$ (vm) sudo apt-get install libopenblas-dev -y
$ (vm) pip install numpy
$ (vm) pip install torch torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html

Verificar que PyTorch puede acceder a las TPUs

Usa el siguiente comando para verificar que PyTorch puede acceder a tus TPUs:

$ (vm) PJRT_DEVICE=TPU python3 -c "import torch_xla.core.xla_model as xm; print(xm.get_xla_supported_devices(\"TPU\"))"

La salida del comando debería tener el siguiente aspecto:

['xla:0', 'xla:1', 'xla:2', 'xla:3', 'xla:4', 'xla:5', 'xla:6', 'xla:7']

Hacer un cálculo básico

Crea un archivo llamado tpu-test.py en el directorio actual y copia y pega el siguiente script en él:

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)

Ejecuta la secuencia de comandos:

(vm)$ PJRT_DEVICE=TPU python3 tpu-test.py

El resultado del script muestra el resultado del cálculo:

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

Limpieza

Para evitar que se apliquen cargos en tu cuenta de Google Cloud por los recursos utilizados en esta página, sigue estos pasos.

Desconéctate de la instancia de TPU de Cloud, si aún no lo has hecho:
```
(vm)$ exit
```
A continuación, se mostrará el mensaje username@projectname, que indica que estás en Cloud Shell.

Elimina tu TPU de Cloud.

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

Para comprobar que los recursos se han eliminado, ejecuta el siguiente comando. Comprueba que tu TPU ya no aparezca en la lista. El proceso de eliminación puede tardar varios minutos.
```
$ gcloud compute tpus tpu-vm list \
    --zone=$ZONE
```

Siguientes pasos

Consulta más información sobre las VMs de TPU de Cloud: