Criação de perfil de desempenho do PyTorch XLA

Visão geral

Veja neste guia 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 a depurar e otimizar o desempenho da carga de trabalho de treinamento.

Conceitos

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

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

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

  1. Crie variáveis para o ID do projeto, o bucket do Cloud Storage e a zona a ser usada para os recursos de 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

  2. Criar uma instância de VM do Compute Engine. É nele que todos os seus scripts e modelos 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

  3. Crie 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

  4. Criar um bucket do Cloud Storage

    Primeiro, instale a CLI gsutil se ela ainda não estiver instalada: instruções de instalação.

  5. Use o comando gsutil mb para criar um bucket do Cloud Storage em que todos os artefatos de criação de perfil estão armazenados. Substitua as variáveis de região e de nome do bucket pelos valores que você usará para o treinamento.

    gsutil mb -p ${PROJECT_ID} -c standard -l REGION gs://${BUCKET_NAME}

    onde:

    • REGION é a região onde você criou o Cloud TPU, por exemplo, europe-west4.

  6. 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

  7. Exporte a conta de serviço e conceda permissões de conta de serviço 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

  1. ssh para a porta de encaminhamento 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"

  2. Crie um ambiente conda dedicado à instalação do TensorBoard:

    conda create -y -n tensorboard python=3.6
conda activate tensorboard
pip install tf-nightly==2.5.0.dev20210211 tb-nightly==2.5.0a20210211 tbp-nightly==2.5.0a20210211

  3. Teste sua instalação executando o servidor TensorBoard em sua VM do Compute Engine e, em seguida, tentando se conectar ao servidor visitando http://localhost:9001/#profile em 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 visitar http://localhost:9001/#profile em sua máquina local, você verá algo semelhante a isto:

imagem

Se você visitar http://localhost:9001, também será possível acessar a página do perfil acima selecionando a opção PROFILE na lista suspensa no canto superior direito, ao lado do botão "UPLOAD" :

imagem

Gerar um perfil do modelo

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

  1. Na nova janela do terminal, exporte novamente o ID do projeto, o ambiente do bucket de armazenamento e as variáveis de zona, já que isso está em um novo shell.

    export PROJECT_ID=PROJECT_ID
export ZONE=ZONE

  2. 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"

  3. Verifique se os testes de integração estão funcionando de ponta a ponta no seu ambiente:

    python /usr/share/torch-xla-1.8/pytorch/xla/test/test_profiler.py # takes <1 min
>

  4. Antes de começar a treinar, edite as linhas a seguir 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 os gráficos dinâmicos e as buscas de tensor, que serão explorados posteriormente na seção de análise de métricas automáticas. Por enquanto, como você está apenas criando o perfil da TPU, o exemplo a seguir é executado com desempenho nominal.

  5. Inicie uma execução de treinamento que seja usada na criação de 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 de TPU (servidor)

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

imagem

Esta página é recarregada automaticamente:

imagem

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

  • Visão geral (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 de nível do framework, estatísticas do PyTorch)
  • Visualizador de traces (requer 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 operação HL do 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 lote de rótulos enviados) Para ver o uso de memória de um módulo HLO específico, selecione o módulo na lista suspensa Hosts à esquerda:

imagem

Depois de visualizar o módulo HLO selecionado, você tem uma visão holística do cronograma de ocupação do HBM de execução do módulo. Isso é determinado pelo tamanho da alocação, pelo tamanho de execução do programa e pelo tamanho do preenchimento.

imagem

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

imagem

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, 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)
   ...

Além disso, a partir do namespace da operação test_mnist, é possível saber que esse módulo HLO corresponde ao loop de avaliação e tem o gerenciador de contexto xp.trace('test_mnist').

Perfil de operação do XLA

O perfil de operações é uma ferramenta do Cloud TPU que exibe 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, com a porcentagem do tempo gasto em operações por categoria e de utilização de FLOPS de TPU;
  • as operações mais demoradas, que são alvos potenciais de otimização;
  • detalhes de operações individuais, incluindo forma, preenchimento e expressões que usam a operação.

É possível usar o perfil de operações 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.

imagem

Descrição de cada um:

  1. Visão geral. Mostra a utilização do Cloud TPU e fornece sugestões para otimização.
  2. 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 são classificadas.
  3. Tabela de operações. 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 do uso do Cloud TPU.
  4. Cards de detalhes das operações. Detalhes sobre a operação exibidos quando você passa o cursor sobre uma operação da tabela. Eles incluem a utilização de FLOPS, a expressão em que a operação é usada e o layout da operação (ajuste).

Visualizador de pod

Para ver uma descrição completa da ferramenta do visualizador de pods, consulte Ferramentas de TPU.

Estatísticas de framework (estatísticas do Tensorflow/PyTorch)

As estatísticas do framework fornecem uma descrição detalhada das estatísticas de operação PyTorch e XRT em execução em dispositivos TPU 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 perfil de eventos de trace do Chrome. Portanto, é necessário usar 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 PyTorch / XLA normalmente, a máquina host funciona principalmente na alocação e alocação de buffer e compilação, enquanto a TPU executa o treinamento do modelo de fato.
  • Com o visualizador de traces, é possível identificar problemas de desempenho do modelo e tomar providências para solucioná-los. Em uma análise mais detalhada, é possível identificar quais operações PyTorch / XLA estão demorando mais para serem executadas.

Observe que é possível adicionar traces diretamente para medir quanto tempo determinadas partes do modelo demoram para ser executadas adicionando anotações xp.Trace(NAME). Por exemplo, o trace a seguir mostra:

imagem

  1. Gerado por anotações explícitas do usuário no código do modelo test_profile_mp_mnist.py.
  2. Operações PyTorch executadas (antes da redução)
  3. Nome do módulo HLO gerado automaticamente pelo PyTorch / XLA.
  4. Operações do XLA executadas no dispositivo (fundado).

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

Criação de perfil de cliente PyTorch / XLA

Assim como quando você cria o perfil do lado da TPU enquanto a execução do modelo estava em andamento, agora você criará o perfil do lado do cliente do 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, veja, que você pode consultar do TensorBoard para capturar um trace.

imagem

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, fornecer uma lista de localhost:port concatenadas por ",". O TensorBoard não é compatível com a visualização de traces de vários processos ao mesmo tempo. Portanto, você verá listas suspensas Host diferentes para cada processo.

O diagrama a seguir mostra um exemplo de trace:

imagem

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

Os traces têm metadados adicionais que podem ser inspecionados. Por exemplo, os traces do TransferToServer e do TransferTransferServer, mostram o número exato de tensores enviados e recebidos e seu tamanho total:

imagem

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

imagem

Além disso, em vez de criar um perfil por meio da IU do TensorBoard, também fornecemos uma API para criar o perfil de maneira programática da TPU e do cliente do PyTorch / XLA: xp.trace().

Análise de métricas automáticas

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áfico muito lenta
  • Transferências frequentes do XLA→CPU
  • HBM repetido do dispositivo para hospedar troca de RAM
  • Desfragmentação HBM repetida
  • aten:: operações 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 compilações artificiais e buscas de tensor. dynamic_graph=True altera artificialmente o tamanho do lote de cada etapa, fazendo com que os gráficos reduzidos do XLA sejam diferentes em cada etapa e recompilação. fetch_often=True insere chamadas loss.item() em cada etapa, resultando na busca de valores de tensores do dispositivo em cada etapa, diminuindo a velocidade do 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 é executar com --num_cores=1, porque simplifica o processo de depuração. Alguns dos resultados de amostra são:

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 de 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 ciclo de treinamento que recupera o valor de loss.item() em cada etapa. Além disso, você poderá ver um aviso CompileTime muito frequente quando o modelo precisar ser recompilado repetidamente devido às 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()

Em seguida, pressione Ctrl^C para fora do script de treinamento e, no final, você verá um resumo das operações reduzidas. Observe que aten::_local_scalar_dense é uma operação especial que corresponde à recuperação de tensores XLA para o contexto da CPU.

Neste relatório, você vê que há dois locais principais em que a operação aten::_local_scalar_dense está sendo chamada, ambas 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 (de acordo com a 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 PT_XLA_DEBUG=1 precisa ter a seguinte aparência:

…
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 11a 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 que você vê o relatório acima. Também é possível ver alguns avisos sobre TransferFromServerTime, porque os tensores estão inicialmente no XLA:TPU antes da execução. Porém, como a operação não é reduzida, primeiro é necessário transferir os tensores do XLA de volta para CPU, execute a operação aten:: na CPU e transfira-a de volta.

Se você quiser 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, em seguida, exclua o bucket da TPU, da VM e do Cloud Storage executando os seguintes comandos:

(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 do TPU + criação de perfil PyTorch/XLA

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

Exportar variáveis de ambiente

  1. Crie variáveis para o ID do projeto e a zona a serem usadas para os recursos de TPU.

    export PROJECT_ID=PROJECT_ID
export ZONE=ZONE

Criar uma Cloud TPU

Consulte a Guia do usuário da VM do TPU e, após a configuração, crie uma VM TPU v3-8, que vem comtorch ,torch_xla, torchvision etensorboard pré-instalado.

  1. Crie 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

  1. Conecte-se por SSH à VM, instale o tensorboard-plugin-profile e inicie um servidor 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:

imagem

Se você visualizar a saída do TensorBoard em http://localhost:9001, também é possível acessar a página do perfil acima selecionando a opção PROFILE na lista suspensa, no canto superior direito, ao lado do botão "UPLOAD":

imagem

Gerar um perfil do modelo

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

  1. Na nova janela do terminal, exporte as variáveis de ID e projeto de projeto novamente, já que isso está em um novo shell.

    export PROJECT_ID=PROJECT_ID
export ZONE=ZONE

  2. ssh na VM:

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

  3. Clone o repositório PyTorch/XLA e execute nosso 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
>

  4. 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 os gráficos dinâmicos e as buscas de tensor, que serão explorados posteriormente na seção de análise de métricas automáticas. Por enquanto, como você está apenas criando o perfil da TPU, o exemplo a seguir é executado com desempenho nominal.

  5. Comece um treinamento:

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

TPU + criação de perfil de cliente

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

imagem

Você verá a seguinte página recarregada:

imagem

Atualmente, na configuração da VM da TPU, somente a ferramenta do visualizador de traces está selecionada. Para isso, no menu suspenso Ferramentas, selecione trace_viewer e inspecione os traces. Na configuração da VM do TPU, veja os traces do lado do cliente e do dispositivo em uma visualização completa:

imagem

Limpeza

  1. Saia da VM e, em seguida, exclua o bucket da TPU, da VM e do Cloud Storage executando os seguintes comandos:

    (vm)$ exit

Exclua a VM da TPU criada:

  1. 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}

  2. Verifique se os recursos foram excluídos executando o comando a seguir. 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.