Como treinar modelos PyTorch no Cloud TPU

Neste documento, descrevemos em detalhes como executar modelos PyTorch em Cloud TPUs. A maneira como você treina modelos do TensorFlow no Cloud TPU depende da arquitetura do Cloud TPU que você está usando. Para saber mais, consulte Arquitetura do sistema.

VMs do Cloud TPU

Com as VMs do Cloud TPU, não há VM do usuário e o modelo e o código são executados em uma VM na máquina host do TPU. É possível usar o SSH diretamente no host da TPU. É possível executar pacotes de instalação de código arbitrários, ver registros, depurar código etc.

O PyTorch/XLA cria um servidor local do TensorFlow sempre que um novo programa PyTorch/XLA é executado. O cliente XRT (XLA Runtime) se conecta ao servidor local do TensorFlow.

image

Para saber mais sobre o treinamento de um modelo PyTorch em VMs de TPU, consulte o Guia de início rápido do PyTorch da VM do Cloud TPU.

Nós legados de TPU

Ao usar nós legados de TPU, você precisa de uma VM de usuário para cada TPU. As VMs do usuário se comunicam com os nós da TPU por meio de chamadas gRPC. O nó de TPU executa um servidor do TensorFlow. O cliente XRT (XLA Runtime) se conecta ao servidor do TensorFlow pela rede.

image