Guia de início rápido do PyTorch da VM do Cloud TPU

Visão geral: este guia de início rápido apresenta uma breve introdução de como trabalhar com o PyTorch e o Cloud TPU.

Antes de começar

Antes de iniciar um dos guias de início rápido de VM do Cloud TPU, leia o Guia do usuário do Cloud TPU. Descreve como configurar um projeto do Google Cloud e como criar uma VM do Cloud TPU.

Faça login na sua Conta do Google. Se você ainda não tiver uma, inscreva-se em uma nova conta. No Console do Google Cloud, selecione ou crie um projeto do Cloud na página do seletor de projetos. Verifique se o faturamento está ativado no projeto.

Ativar a API Cloud TPU

Ative a API Cloud TPU usando o seguinte comando gcloud no Cloud Shell. Também é possível ativá-lo no Console do Google Cloud.

$ gcloud services enable tpu.googleapis.com

Configurar o comando gcloud

Execute os comandos a seguir para configurar gcloud para usar o projeto do GCP e instalar componentes necessários para a visualização da VM do TPU.

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

Crie uma VM do Cloud TPU com gcloud

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

Campos obrigatórios

zone
A zona em que você planeja criar a Cloud TPU.
accelerator-type
O tipo do Cloud TPU a ser criado.
version
A versão do ambiente de execução do Cloud TPU.

Conecte-se à VM do Cloud TPU

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

Campos obrigatórios

tpu_name
O nome da VM da TPU a que você está se conectando.
zone
A zona em que você está criando o Cloud TPU. Atualmente europe-west4-a e us-central1-a são compatíveis.
project-id
O ID do projeto do GCP.

Definir configuração de dispositivo da TPU XRT

$ export XRT_TPU_CONFIG="localservice;0;localhost:51011"

Realize um cálculo simples:

Inicie o interpretador do Python 3:

(vm)$ python3

Saída do 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.

Importar pacotes do PyTorch

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

Realizar o cálculo

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

Saída do comando:

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

Saia do interpretador do Python:

>>> exit()

Limpeza

Para evitar cobranças na sua conta do Google Cloud pelos recursos usados neste guia de início rápido, siga estas etapas:

  1. Desconecte-se da instância do Compute Engine, caso ainda não tenha feito isso:

    (vm)$ exit

    Agora, o prompt será username@projectname, mostrando que você está no Cloud Shell.

  2. Exclua o Cloud TPU.

    $ gcloud alpha compute tpus tpu-vm delete tpu-name \
      --zone=us-central1-b
  3. Execute gcloud alpha compute tpus tpu-vm list para verificar se os recursos foram excluídos. A exclusão pode levar vários minutos. Uma resposta como esta indica que suas instâncias foram excluídas com sucesso.

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

     NAME             STATUS

A seguir

Leia mais sobre VMs do Cloud TPU: