Ejecuta un cálculo en una VM de Cloud TPU con PyTorch
En este documento, se proporciona una breve introducción al trabajo con PyTorch y Cloud TPU.
Antes de comenzar
Antes de ejecutar los comandos de este documento, debes crear una cuenta de Google Cloud , instalar Google Cloud CLI y configurar el comando gcloud
. Para obtener más información, consulta Configura el entorno de Cloud TPU.
Crea una Cloud TPU con gcloud
Define algunas variables de entorno para facilitar el uso de los comandos.
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 Google Cloud ID de tu proyecto. Usa un proyecto existente o crea uno nuevo. TPU_NAME
El nombre de la TPU. ZONE
Es 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 Cloud TPU que deseas 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 Cloud TPU Ejecuta el siguiente comando para crear tu VM de TPU:
$ 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 Cloud TPU
Conéctate a tu VM de TPU a través de 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, es posible que sea porque 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 Cómo conectarse a una VM de TPU sin una dirección IP pública.
Instala 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
Verifica que PyTorch pueda acceder a las TPU
Usa el siguiente comando para verificar que PyTorch pueda acceder a tus TPU:
$ (vm) PJRT_DEVICE=TPU python3 -c "import torch_xla.core.xla_model as xm; print(xm.get_xla_supported_devices(\"TPU\"))"
El resultado del comando debería verse de la siguiente manera:
['xla:0', 'xla:1', 'xla:2', 'xla:3', 'xla:4', 'xla:5', 'xla:6', 'xla:7']
Realiza un cálculo básico
Crea un archivo llamado
tpu-test.py
en el directorio actual y copia y pega la siguiente secuencia de comandos 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 de la secuencia de comandos muestra el resultado del procesamiento:
tensor([[-0.2121, 1.5589, -0.6951], [-0.7886, -0.2022, 0.9242], [ 0.8555, -1.8698, 1.4333]], device='xla:1')
Limpia
Para evitar que se apliquen cargos a tu cuenta de Google Cloud por los recursos que usaste en esta página, sigue estos pasos.
Desconéctate de la instancia de Cloud TPU, si aún no lo hiciste:
(vm)$ exit
El mensaje ahora debería mostrar
username@projectname
, que indica que estás en Cloud Shell.Borra tu Cloud TPU.
$ gcloud compute tpus tpu-vm delete $TPU_NAME \ --project=$PROJECT_ID \ --zone=$ZONE
Ejecuta el siguiente comando para verificar que los recursos se hayan borrado. Asegúrate de que tu TPU ya no aparezca en la lista. La eliminación puede tardar varios minutos.
$ gcloud compute tpus tpu-vm list \ --zone=$ZONE
¿Qué sigue?
Obtén más información sobre las VMs de Cloud TPU: