Esta página foi traduzida pela API Cloud Translation.

Criação de perfis de desempenho do PyTorch XLA

Visão geral

Neste guia, você verá como usar as ferramentas de desempenho do Cloud TPU e o recurso de análise automática de métricas com o PyTorch. Essas ferramentas ajudam você a depurar e otimizar o desempenho da carga de trabalho de treinamento.

Para mais informações sobre o desempenho do Cloud TPU com o PyTorch, consulte as seguintes postagens do blog:

Como escalonar cargas de trabalho de aprendizado profundo PyTorch
Depuração de desempenho do PyTorch XLA em VMs de TPU: parte 1
Depuração de desempenho do PyTorch XLA em VMs de TPU: parte 2
Depuração de desempenho do PyTorch XLA em VMs de TPU: parte 3
[Desempenho lento do tensor com PyTorch XLA](https://pytorch.org/blog/understanding-lazytensor-system-performance-with-pytorch-xla-on-cloud-tpu/

Se você for novo no PyTorch/XLA, consulte os documentos API_GUIDE e solução de problemas do PyTorch. Para o Cloud TPU, consulte o documento de conceitos.

Nó da TPU + criação de perfil PyTorch/XLA

Criar recursos relacionados ao Cloud TPU

Criar e inicializar os recursos necessários relacionados ao Cloud TPU.

Criar variáveis para o ID do projeto, seu bucket do Cloud Storage e a zona a ser usada nos recursos da TPU.

export PROJECT_ID=PROJECT_ID
export BUCKET_NAME=BUCKET_NAME
export ZONE=ZONE

gcloud --project=$PROJECT_ID compute project-info add-metadata \
--metadata BUCKET_NAME=$BUCKET_NAME

Criar uma instância de VM do Compute Engine. É onde todos os modelos e scripts do Python são armazenados.

gcloud compute instances create profiler-tutorial-vm \
  --project=${PROJECT_ID} \
  --zone=${ZONE} \
  --machine-type=n1-standard-16 \
  --image-project=ml-images \
  --image-family=torch-xla \
  --boot-disk-size=300GB \
  --scopes=https://www.googleapis.com/auth/cloud-platform

Criar um recurso de TPU.

gcloud compute tpus create profiler-tutorial-tpu \
 --project=${PROJECT_ID} \
 --zone=${ZONE} \
 --network=default \
 --version=pytorch-1.8 \
 --accelerator-type=v3-8

Criar um bucket do Cloud Storage

Primeiro, instale a CLI gsutil se ainda não a tiver instalado: instruções de instalação.
Use o comando gsutil mb para criar um bucket do Cloud Storage em que todos os artefatos de criação de perfil são armazenados. Substitua as variáveis de região e nome de bucket pelos valores que você usará no treinamento.
```
gsutil mb -p ${PROJECT_ID} -c standard -l REGION gs://${BUCKET_NAME}
```
em que:
- REGION é a região em que você criou o Cloud TPU, por exemplo, europe-west4.
Crie uma conta de serviço para o projeto da Cloud TPU.
```
gcloud beta services identity create --service tpu.googleapis.com --project $PROJECT_ID
```
O comando retorna uma conta de serviço do Cloud TPU com o formato a seguir:
```
service-PROJECT_NUMBER@cloud-tpu.iam.gserviceaccount.com
```
Por exemplo: service-164006649440@cloud-tpu.iam.gserviceaccount.com

Exporte a conta de serviço e conceda permissões a ela no bucket de armazenamento. Substitua account-number pelo PROJECT_NUMBER retornado na saída da criação da conta de serviço.

export SERVICE_ACCOUNT=service-ACCOUNT_NUMBER@cloud-tpu.iam.gserviceaccount.com
gsutil acl ch -u $SERVICE_ACCOUNT:READER gs://${BUCKET_NAME}
gsutil acl ch -u $SERVICE_ACCOUNT:WRITER gs://${BUCKET_NAME}

Configurar o TensorBoard

ssh para a porta de encaminhamento de VMs 9001 na VM para a porta 9001 na máquina local. Essa porta é usada para abrir a IU do TensorBoard no navegador local.
```
  gcloud compute ssh profiler-tutorial-vm \
   --project ${PROJECT_ID} \
   --zone ${ZONE} \
   --ssh-flag="-4 -L 9001:localhost:9001"
```

Criar um ambiente conda dedicado à instalação do TensorBoard:

conda create -y -n tensorboard python=3.6
conda activate tensorboard
pip install tf-nightly==2.6.0.dev20210511 tb-nightly==2.6.0a20210512 tbp-nightly==2.5.0a20210511

Para testar sua instalação, execute o servidor do TensorBoard na VM do Compute Engine e tente se conectar ao servidor em http://localhost:9001/#profile na sua máquina local:

 # Get bucket name
 BUCKET_NAME=$(curl "http://metadata.google.internal/computeMetadata/v1/project/attributes/BUCKET_NAME" -H "Metadata-Flavor: Google")

 tensorboard --logdir gs://${BUCKET_NAME} --port 9001

Ao acessar http://localhost:9001/#profile na máquina local, você verá algo parecido com o seguinte:

Se você acessar http://localhost:9001, também poderá acessar a página de perfil acima. Basta selecionar a opção PERFIL na lista suspensa no canto superior direito ao lado do botão ENVIAR:

Criar o perfil do modelo

Para manter o servidor do TensorBoard ativo, na máquina local, inicie uma nova janela de terminal e ssh na VM do GCE novamente, (desta vez sem usar a opção -L de encaminhamento de portas).

Na nova janela do terminal, exporte novamente o ID do projeto, o ambiente do bucket de armazenamento e as variáveis de zona, porque ele está em um novo shell.
```
export PROJECT_ID=PROJECT_ID
export ZONE=ZONE
```

ssh na VM:

  gcloud compute ssh profiler-tutorial-vm \
   --project ${PROJECT_ID} \
   --zone ${ZONE}

conda activate torch-xla-1.8
PROJECT_ID=$(curl "http://metadata.google.internal/computeMetadata/v1/project/project-id" -H "Metadata-Flavor: Google")
export TPU_IP_ADDRESS=$(gcloud compute tpus describe profiler-tutorial-tpu --zone=${ZONE} --project=${PROJECT_ID} \
--format="value(ipAddress)")
echo TPU_IP_ADDRESS=${TPU_IP_ADDRESS}
export XRT_TPU_CONFIG="tpu_worker;0;$TPU_IP_ADDRESS:8470"

Verifique se os testes de integração estão funcionando de ponta a ponta no ambiente:
```
python /usr/share/torch-xla-1.8/pytorch/xla/test/test_profiler.py # takes <1 min
>
```
Antes de começar a treinar, edite as seguintes linhas em /usr/share/torch-xla-1.8/pytorch/xla/test/test_profile_mp_mnist.py:

Mudar:
```
accuracy = train_mnist(flags, dynamic_graph=True, fetch_often=True)
```
Para:
```
accuracy = train_mnist(flags, dynamic_graph=False, fetch_often=False)
```
Os dois argumentos acima em train_mnist causam artificialmente gráficos dinâmicos e buscas de tensores, que serão explorados posteriormente na seção "Análise de métricas automáticas". Por enquanto, como você está apenas analisando a TPU, o exemplo a seguir é executado com desempenho nominal.

Inicie uma execução de treinamento usada para criar o perfil do servidor:

PT_XLA_DEBUG=1 XLA_HLO_DEBUG=1 python /usr/share/torch-xla-1.8/pytorch/xla/test/test_profile_mp_mnist.py --num_epochs 1000 --fake_data

Criação de perfil da TPU (servidor)

Quando o treinamento estiver em execução, acesse http://localhost:9001/#profile e capture um perfil usando as seguintes instruções:

A página a seguir será recarregada automaticamente:

As ferramentas compatíveis são mostradas na lista suspensa Ferramentas no painel esquerdo:

Página de visão geral (isso não inclui a exibição do analisador de pipeline de entrada, que é fundamentalmente diferente entre o TensorFlow/TPU e a PyTorch/XLA TPU).
Visualizador de memória
Perfil das operações
Visualizador de pod
Estatísticas do TensorFlow (estatísticas no nível do framework, ou seja, estatísticas do PyTorch)
Visualizador do Trace (exige o navegador Chrome)

Página de visão geral

Esta página mostra uma visão geral de um perfil capturado. Este exemplo mostra um tempo de inatividade muito alto porque está treinando um pequeno modelo no conjunto de dados MNIST.

Visualizador de memória

Mostra a memória do dispositivo (HBM) usada por tensor e op HLO PyTorch. O visualizador de memória captura uma visualização da memória por módulo HLO. Portanto, haverá alguns módulos como gráficos de alocação de dados do dispositivo (entrada e identificadores enviados em lote). Para ver o uso de memória de um módulo HLO específico, selecione o módulo na lista suspensa Hosts à esquerda:

Ao visualizar o módulo HLO selecionado, você tem uma visão holística da linha do tempo de ocupação da HBM de execução do módulo. Isso é ordenado por tamanho de alocação, tamanho de execução do programa e tamanho do padding.

Cada uma das alocações de buffer pode ser inspecionada ainda passando o cursor sobre elas. Por exemplo, para informar qual alocação está ocupando a maior parte da HBM do dispositivo:

No exemplo acima, (1) corresponde às anotações torch_xla.debug.profiler.Trace adicionadas pelo código do usuário. Ao inspecionar o código test_profile_mp_mnist.py, ele corresponde a esta linha:

   class MNIST(nn.Module):
   ...
     def forward(self, x):
       with xp.Trace('conv1'):
         x = F.relu(F.max_pool2d(self.conv1(x), 2))
         x = self.bn1(x)
   ...

Também no namespace op test_mnist, é possível ver que esse módulo HLO corresponde ao loop de avaliação porque tem o gerenciador de contexto xp.trace('test_mnist').

Perfil de op do XLA

O perfil de op é uma ferramenta do Cloud TPU que exibe as estatísticas de desempenho das operações do XLA executadas durante um período de criação de perfil. O perfil de operações mostra:

O uso que o aplicativo faz do Cloud TPU como uma porcentagem do tempo gasto em operações por categoria e de utilização de FLOPS de TPU;
as operações mais demoradas, Essas operações são alvos potenciais de otimização.
Detalhes de operações individuais, incluindo forma, padding e expressões que usam a operação.

É possível usar o perfil de op para identificar bons alvos de otimização. Por exemplo, se o modelo atingir apenas 5% de FLOPS de pico da TPU, será possível usar a ferramenta para identificar quais operações do XLA são mais demoradas na execução e quantos FLOPS de TPU elas consomem.

Descrição de cada item:

Seção de visão geral. Mostra a utilização do Cloud TPU e fornece sugestões para otimização.
Painel de controle. Contém controles que permitem definir o número de operações exibidas na tabela, quais operações são exibidas e como elas são classificadas.
Tabela op. Uma tabela que lista as principais categorias de operações do TensorFlow associadas às operações do XLA. Essas operações são classificadas por porcentagem de uso do Cloud TPU.
Cartões de detalhes de operações Detalhes sobre a operação exibidos quando você passa o cursor sobre uma operação da tabela. Inclui a utilização de FLOPS, a expressão em que a operação é usada e o layout da operação (fit).

Visualizador de pod

Consulte Ferramentas da TPU para uma descrição completa da ferramenta de visualização de pods.

Estatísticas do framework (estatísticas do TensorFlow/PyTorch)

As estatísticas do framework fornecem uma análise detalhada das estatísticas de operações PyTorch e XRT em execução em dispositivos e hosts de TPU.

Visualizador de traces

O visualizador de traces é uma ferramenta de análise de desempenho do Cloud TPU. O visualizador de traces usa o visualizador de criação de perfis de eventos de trace do Chrome para que use o navegador Chrome.

O visualizador de traces contém uma linha do tempo que mostra:

As durações das operações que foram executadas pelo seu modelo do TensorFlow.
qual parte do sistema (TPU ou máquina host) executou uma operação. Para o PyTorch/XLA geralmente, a máquina host trabalha principalmente na compilação e alocação/desalocação de buffer, enquanto a TPU executa o treinamento do modelo real.
Com o visualizador de traces, é possível identificar problemas de desempenho do modelo e tomar providências para solucioná-los. Detalhando, é possível identificar quais operações PyTorch/XLA levam mais tempo para serem executadas.

É possível adicionar diretamente traces para medir quanto tempo determinadas partes do modelo levam para serem executadas. Para isso, adicione anotações do xp.Trace(NAME). Por exemplo, o trace a seguir é exibido:

Gerado por anotações de usuário explícitas no código de modelo de test_profile_mp_mnist.py.
Operações do PyTorch executadas (pré-redução).
Nome do módulo HLO gerado automaticamente pelo PyTorch/XLA.
Operações de XLA executadas no dispositivo (fundido).

Para informações mais detalhadas, consulte a documentação genérica da TPU para o visualizador de traces, mas ignore as seções sobre pipeline de entrada e outras partes específicas do TensorFlow, já que elas não são relevantes no contexto deste documento.

Criação de perfil do cliente PyTorch/XLA

Assim como na criação de perfil da TPU enquanto a execução do modelo estava em andamento, agora você criará o perfil do lado do cliente PyTorch/XLA durante o treinamento. A principal ferramenta de monitoramento usada no lado do cliente é o Visualizador de traces.

É preciso iniciar o servidor de criação de perfil no script de treinamento. Por exemplo, consulte, que pode ser consultado no TensorBoard para capturar um trace.

Para capturar traces de vários processos, cada processo pode iniciar servidores de criação de perfil em portas diferentes (por exemplo, adicionando xm.get_ordinal() a um número de porta base) e, em seguida, fornecendo uma lista de localhost:port concatenado por ','. Como o TensorBoard não é compatível com a visualização de traces de vários processos ao mesmo tempo, você verá menus suspensos de Host diferentes para cada um deles.

O diagrama a seguir mostra um trace de amostra:

Da mesma forma que diferentes traces de namespace podem ser adicionados a traces de TPU, é possível usar a mesma API para adicioná-los aos traces do lado do cliente (xp.Trace(NAME)). Como esse modelo é pequeno e está sendo usado com pequenas imagens MNIST, os tempos das etapas serão curtos e não necessariamente uniformes. Como exercício, você pode tentar adicionar os traces e inicializar um servidor do criador de perfil semelhante ao do nosso exemplo ao test_train_mp_imagenet.py --fake_data para receber traces de ResNet50.

Os traces têm metadados adicionais que podem ser inspecionados. Por exemplo, os traces TransferToServer e TransferFromServer mostram o número exato de tensores que estão sendo enviados e recebidos e o tamanho total deles:

Para compilações de gráficos XLA, é possível ver o hash do gráfico que pode ser útil para diagnosticar problemas:

Além disso, em vez de criar perfis da IU do TensorBoard, também fornecemos a API para a criação de perfil programática da TPU e do cliente do PyTorch/XLA: xp.trace().

Análise automática de métricas

Nesta seção, você verá como usar o modo de depuração para detectar problemas de desempenho, como:

Gráficos dinâmicos/compilações contínuas
Compilação de gráficos muito lenta
Execução de gráficos muito lenta
Transferências frequentes XLA→CPU
HBM do dispositivo repetido para hospedar a troca de RAM
Desfragmentação de HBM repetida
Operações aten:: não reduzidas

Antes de iniciar o treinamento, reverta as seguintes linhas em /usr/share/torch-xla-1.8/pytorch/xla/test/test_profile_mp_mnist.py:

Mudar:

   accuracy = train_mnist(flags, dynamic_graph=False, fetch_often=False)

Para:

   accuracy = train_mnist(flags, dynamic_graph=True, fetch_often=True)

Essas alterações causam artificialmente compilações e buscas de tensor. dynamic_graph=True muda artificialmente o tamanho do lote para cada etapa, fazendo com que os gráficos XLA diminuídos sejam diferentes em cada etapa e recompilação. fetch_often=True insere chamadas loss.item() em todas as etapas, resultando na busca de valores de tensor do dispositivo em cada etapa, reduzindo o desempenho.

Como executar um exemplo de script de treinamento:

   PT_XLA_DEBUG=1 python /usr/share/torch-xla-1.8/pytorch/xla/test/test_profile_mp_mnist.py --fake_data --num_cores=1

Durante a depuração, a prática recomendada é executá-la com --num_cores=1, pois ele simplifica o processo de depuração. Alguns dos resultados de amostra são assim:

Epoch 1 train begin 01:18:05
| Training Device=xla:1/0 Step=0 Loss=0.00000 Rate=1905.00 GlobalRate=1904.72 Time=01:18:05
pt-xla-profiler: TransferFromServerTime too frequent: 3 counts during 3 steps
pt-xla-profiler: TransferFromServerTime too frequent: 4 counts during 4 steps
pt-xla-profiler: TransferFromServerTime too frequent: 5 counts during 5 steps
pt-xla-profiler: TransferFromServerTime too frequent: 6 counts during 6 steps
pt-xla-profiler: TransferFromServerTime too frequent: 7 counts during 7 steps
pt-xla-profiler: TransferFromServerTime too frequent: 8 counts during 8 steps
pt-xla-profiler: TransferFromServerTime too frequent: 9 counts during 9 steps
pt-xla-profiler: TransferFromServerTime too frequent: 10 counts during 10 steps
pt-xla-profiler: CompileTime too frequent: 21 counts during 11 steps
pt-xla-profiler: TransferFromServerTime too frequent: 11 counts during 11 steps
pt-xla-profiler: CompileTime too frequent: 23 counts during 12 steps

As linhas com o prefixo pt-xla-profiler correspondem à saída da análise de métricas automáticas. Neste exemplo, é possível ver que TransferFromServerTime foi visto com muita frequência, uma vez por etapa. Isso se deve ao loop de treinamento que recupera o valor de loss.item() em todas as etapas. Além disso, você poderá ver um aviso CompileTime de forma muito frequente, quando o modelo precisar ser recompilado repetidamente devido a formas dinâmicas no gráfico. O snippet de código a seguir causa esse tipo de problema: test_profile_mp_mnist.py

    for step, (data, target) in enumerate(loader):
      if dynamic_graph:
        # The batch dimension is different every step.
        index = max(-step, -flags.batch_size + 1)  # non-empty
        data, target = data[:-index, :, :, :], target[:-index]
      ...
      if fetch_often:
        # Fetch tensor value from XLA:TPU to CPU every step.
        loss_i = loss.item()

Depois, pressione Ctrl^C para sair do script de treinamento. Quando terminar, você verá um resumo das operações não reduzidas. Observe que aten::_local_scalar_dense é uma operação especial que corresponde à recuperação de tensores XLA de volta ao contexto da CPU.

Neste relatório, você vê que há dois lugares principais em que a operação aten::_local_scalar_dense está sendo chamada, ambos correspondem ao código-fonte de loss.item():

test/test_profile_mp_mnist.py:158
test/test_profile_mp_mnist.py:61

pt-xla-profiler: ================================================================================
pt-xla-profiler: Unlowered Op usage summary (more of these ops, lower performance)
pt-xla-profiler: Note: _local_scalar_dense typically indicates CPU context access
pt-xla-profiler: --------------------------------------------------------------------------------
pt-xla-profiler: FRAME (count=27):
pt-xla-profiler: Unlowered Op: "_local_scalar_dense"
pt-xla-profiler: Python Frames:
pt-xla-profiler:   train_loop_fn (test/test_profile_mp_mnist.py:158)
pt-xla-profiler:   train_mnist (test/test_profile_mp_mnist.py:184)
pt-xla-profiler:   _mp_fn (test/test_profile_mp_mnist.py:206)
pt-xla-profiler:   _start_fn (/home/jysohn/git/jysohn23/pytorch/xla/torch_xla/distributed/xla_multiprocessing.py:323)
pt-xla-profiler:   spawn (/home/jysohn/git/jysohn23/pytorch/xla/torch_xla/distributed/xla_multiprocessing.py:386)
pt-xla-profiler:    (test/test_profile_mp_mnist.py:216)
pt-xla-profiler:
pt-xla-profiler:
pt-xla-profiler: FRAME (count=2):
pt-xla-profiler: Unlowered Op: "_local_scalar_dense"
pt-xla-profiler: Python Frames:
pt-xla-profiler:   _train_update (test/test_profile_mp_mnist.py:61)
pt-xla-profiler:    (/home/jysohn/git/jysohn23/pytorch/xla/torch_xla/core/xla_model.py:700)
pt-xla-profiler:   _run_step_closures (/home/jysohn/git/jysohn23/pytorch/xla/torch_xla/core/xla_model.py:709)
pt-xla-profiler:   mark_step (/home/jysohn/git/jysohn23/pytorch/xla/torch_xla/core/xla_model.py:723)
pt-xla-profiler:   __exit__ (/home/jysohn/git/jysohn23/pytorch/xla/torch_xla/debug/profiler.py:153)
pt-xla-profiler:   train_loop_fn (test/test_profile_mp_mnist.py:162)
pt-xla-profiler:   train_mnist (test/test_profile_mp_mnist.py:184)
pt-xla-profiler:   _mp_fn (test/test_profile_mp_mnist.py:206)
pt-xla-profiler:   _start_fn (/home/jysohn/git/jysohn23/pytorch/xla/torch_xla/distributed/xla_multiprocessing.py:323)
pt-xla-profiler:   spawn (/home/jysohn/git/jysohn23/pytorch/xla/torch_xla/distributed/xla_multiprocessing.py:386)
pt-xla-profiler:    (test/test_profile_mp_mnist.py:216)
pt-xla-profiler:
pt-xla-profiler:
pt-xla-profiler: ================================================================================

Agora, execute a análise de métricas automáticas no script abaixo, que contém uma operação não reduzida (desde nossa versão 1.8), _ctc_loss:

PT_XLA_DEBUG=1 python <<EOF
import torch
import torch_xla.core.xla_model as xm

dev = xm.xla_device()
t = torch.randn(50, 16, 20).log_softmax(2).to(dev)
target = torch.randint(low=1, high=20, size=(16, 30), dtype=torch.long).to(dev)
input_lengths = torch.full(size=(16,), fill_value=50, dtype=torch.long).to(dev)
target_lengths = torch.randint(low=10, high=30, size=(16,), dtype=torch.long).to(dev)

for _ in range(10):
  loss = torch.nn.CTCLoss()(t, target, input_lengths, target_lengths)
  xm.mark_step()
EOF

A execução do script acima com a saída PT_XLA_DEBUG=1 precisa ter uma aparência semelhante à seguinte:

…
pt-xla-profiler: TransferFromServerTime too frequent: 30 counts during 10 steps
pt-xla-profiler: Op(s) not lowered: aten::_ctc_loss,  Please open a GitHub issue with the above op lowering requests.
pt-xla-profiler: ================================================================================
pt-xla-profiler: Unlowered Op usage summary (more of these ops, lower performance)
pt-xla-profiler: Note: _local_scalar_dense typically indicates CPU context access
pt-xla-profiler: --------------------------------------------------------------------------------
pt-xla-profiler: FRAME (count=10):
pt-xla-profiler: Unlowered Op: "_ctc_loss"
pt-xla-profiler: Python Frames:
pt-xla-profiler:   ctc_loss (/anaconda3/envs/torch-xla-1.8/lib/python3.6/site-packages/torch/nn/functional.py:2305)
pt-xla-profiler:   forward (/anaconda3/envs/torch-xla-1.8/lib/python3.6/site-packages/torch/nn/modules/loss.py:1593)
pt-xla-profiler:   _call_impl (/anaconda3/envs/torch-xla-1.8/lib/python3.6/site-packages/torch/nn/modules/module.py:889)
pt-xla-profiler:    (:11)
pt-xla-profiler:
pt-xla-profiler:
pt-xla-profiler: ================================================================================

O analisador de métricas automáticas mostra que a 11ª linha do STDIN está causando essa operação não reduzida (ou seja, a linha com torch.nn.CTCLoss()). Atualmente, a operação ctc_loss não foi reduzida, por isso você vê o relatório acima. Também é possível ver alguns avisos para TransferFromServerTime, já que os tensores estão inicialmente em XLA:TPU antes da execução. Porém, como a operação não é reduzida, primeiro você precisa transferir os tensores XLA de volta para a CPU, executar a operação aten:: na CPU e transferi-la de volta.

Caso você queira gravar a saída pt-xla-profiler em um arquivo, defina PT_XLA_DEBUG=1 e PT_XLA_DEBUG_FILE=$PATH_TO_FILE.

Limpeza

Saia da VM e exclua o bucket da TPU, da VM e do Cloud Storage executando os comandos a seguir:

(vm)$ exit

gcloud compute instances delete profiler-tutorial-vm \
  --zone=${ZONE} \
  --project=${PROJECT_ID}

gcloud compute tpus delete profiler-tutorial-tpu \
  --zone=${ZONE} \
  --project=${PROJECT_ID} \
  --async

gsutil rm -fr gs://${BUCKET_NAME}

VM de TPU + criação de perfil do PyTorch/XLA

Use esta seção para criar o perfil do PyTorch/XLA usando a arquitetura de VM de TPU.

Exportar variáveis de ambiente

Crie variáveis para o ID do projeto e a zona a ser usada nos recursos da TPU.
```
export PROJECT_ID=PROJECT_ID
export ZONE=ZONE
```

Criar uma Cloud TPU

Consulte o Guia do usuário de VMs da TPU e, após a configuração, crie uma VM de TPU v3-8, que vem com torch, torch_xla, torchvision e tensorboard pré-instalados.

Criar um recurso de TPU.

gcloud alpha compute tpus tpu-vm create profiler-tutorial-tpu-vm \
 --project=${PROJECT_ID} \
 --zone=${ZONE} \
 --version=v2-alpha \
 --accelerator-type=v3-8

Inicialização do servidor do TensorBoard

Conecte-se por SSH à VM, instale tensorboard-plugin-profile e inicie um servidor do TensorBoard.

  gcloud alpha compute tpus tpu-vm ssh profiler-tutorial-tpu-vm \
   --project ${PROJECT_ID} \
   --zone ${ZONE} \
   --ssh-flag="-4 -L 9001:localhost:9001"

  pip3 install tf-nightly==2.6.0.dev20210511 tb-nightly==2.6.0a20210512 tbp-nightly==2.5.0a20210511

  tensorboard --logdir ./tensorboard --port 9001

Ao visualizar a saída do TensorBoard em http://localhost:9001 na máquina local, você verá algo assim:

Se você visualizar a saída do TensorBoard em http://localhost:9001, também poderá acessar a página de perfil acima. Basta selecionar a opção PROFILE no menu suspenso no canto superior direito, ao lado do botão UPLOAD:

Criar o perfil do modelo

Em uma nova janela de terminal no ambiente de desenvolvimento, exporte as mesmas variáveis de ambiente mostradas acima e ssh para a VM de TPU:

Na nova janela do terminal, exporte as variáveis de zona e ID do projeto novamente, já que isso está em um novo shell.
```
export PROJECT_ID=PROJECT_ID
export ZONE=ZONE
```

ssh na VM:

  gcloud alpha compute tpus tpu-vm ssh profiler-tutorial-tpu-vm \
   --project ${PROJECT_ID} \
   --zone ${ZONE}

Clone o repositório PyTorch/XLA e execute o teste e2e:

  git clone -b r1.8 https://github.com/pytorch/xla
  export XRT_TPU_CONFIG="localservice;0;localhost:51011"
  python3 xla/test/test_profiler.py  # takes <1 min
>

Antes de iniciar o treinamento, edite as seguintes linhas em xla/test/test_profile_mp_mnist.py:

Mudar:
```
    accuracy = train_mnist(flags, dynamic_graph=True, fetch_often=True)
```
Para:
```
    accuracy = train_mnist(flags, dynamic_graph=False, fetch_often=False)
```
Os dois argumentos acima em train_mnist causam artificialmente gráficos dinâmicos e buscas de tensores, que serão explorados posteriormente na seção "Análise de métricas automáticas". Por enquanto, como você está apenas analisando a TPU, o exemplo a seguir é executado com desempenho nominal.

Inicie uma execução de treinamento:

 XLA_HLO_DEBUG=1 python3 xla/test/test_profile_mp_mnist.py --num_epochs 1000 --fake_data

Criação de perfis de TPU + cliente

Quando o treinamento estiver em execução, veja a saída do TensorBoard em http://localhost:9001 e capture um perfil usando as seguintes instruções:

Você verá a seguinte página recarregada:

Atualmente, na configuração da VM de TPU, somente a ferramenta visualizador de traces está selecionada. Portanto, na lista suspensa Ferramentas, selecione trace_viewer e inspecione os traces. Na configuração da VM de TPU, você vê os traces do lado do "cliente" e do lado do dispositivo da TPU em uma visualização completa:

Limpeza

Saia da VM e exclua o bucket da TPU, da VM e do Cloud Storage executando os comandos a seguir:
```
(vm)$ exit
```

Exclua a VM da TPU que você criou:

Exclua os recursos do Cloud TPU e do Compute Engine.

$ gcloud alpha compute tpus tpu-vm delete profiler-tutorial-tpu-vm \
  --project ${PROJECT_ID} --zone=${ZONE}

Verifique se os recursos foram excluídos executando o seguinte comando. 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 --project ${PROJECT_ID} --zone=${ZONE}
```
```
Listed 0 items.
```