As Unidades de Processamento de Tensor (TPUs, na sigla em inglês) são circuitos integrados específicos para aplicativos (ASICs, na sigla em inglês) desenvolvidos especialmente pelo Google. Elas são usadas para acelerar as cargas de trabalho de machine learning. É possível executar jobs de treinamento
no AI Platform Training com uma Cloud TPU. O AI Platform Training fornece uma
interface de gerenciamento de jobs para que você não precise gerenciar a TPU por conta própria.
Em vez disso, use a API jobs
do AI Platform Training do
mesmo modo que você a usaria para realizar o treinamento em uma CPU ou GPU.
Com as APIs de alto nível do TensorFlow, seus modelos poderão ser executados no hardware da Cloud TPU.
Como configurar o ambiente do Google Cloud
Para configurar o ambiente do Google Cloud, siga as instruções na seção de configuração do guia de primeiros passos.
Como autorizar o Cloud TPU a acessar o projeto
Siga estas etapas para autorizar o nome da conta de serviço do Cloud TPU associada ao seu projeto do Google Cloud:
Chame
projects.getConfig
para receber o nome da conta de serviço do Cloud TPU. Exemplo:PROJECT_ID=PROJECT_ID curl -H "Authorization: Bearer $(gcloud auth print-access-token)" \ https://ml.googleapis.com/v1/projects/$PROJECT_ID:getConfig
Salve os valores dos campos
serviceAccountProject
etpuServiceAccount
retornados pela API.Inicialize a conta de serviço do Cloud TPU:
curl -H "Authorization: Bearer $(gcloud auth print-access-token)" \ -H "Content-Type: application/json" -d '{}' \ https://serviceusage.googleapis.com/v1beta1/projects/<serviceAccountProject>/services/tpu.googleapis.com:generateServiceIdentity
Agora adicione a conta de serviço da Cloud TPU como membro do seu projeto, com o papel Agente de serviço do Cloud ML. Conclua as etapas a seguir no
Console do Google Cloud ou usando o comando gcloud
:
Console
- Faça login no Console do Google Cloud e escolha o projeto em que você usa a TPU.
- Escolha IAM e Admin > IAM.
- Clique no botão Adicionar para adicionar um membro ao projeto.
- Digite a conta de serviço de TPU na caixa de texto Membros.
- Clique na lista suspensa Papéis.
- Ative o papel Agente de serviço do Cloud ML (Agentes de serviço > Agente de serviço do Cloud ML).
gcloud
Defina as variáveis de ambiente que contêm o ID do projeto e a conta de serviço da Cloud TPU:
PROJECT_ID=PROJECT_ID SVC_ACCOUNT=your-tpu-sa-123@your-tpu-sa.google.com.iam.gserviceaccount.com
Atribua o papel
ml.serviceAgent
à conta de serviço do Cloud TPU.gcloud projects add-iam-policy-binding $PROJECT_ID \ --member serviceAccount:$SVC_ACCOUNT --role roles/ml.serviceAgent
Para mais detalhes sobre a concessão de papéis a contas de serviço, consulte a documentação do IAM.
Exemplo: como treinar um modelo MNIST de amostra
Nesta seção, mostramos como treinar um modelo MNIST de amostra usando
uma TPU e a versão 2.11 do ambiente de execução. No exemplo de job, o nível de escalonamento
predefinido BASIC_TPU
é usado para a configuração da máquina.
As seções posteriores do guia contêm informações sobre como definir uma configuração personalizada.
Neste exemplo, pressupomos que você esteja usando um shell do Bash com a CLI gcloud instalada. Execute os seguintes comandos para receber o código e enviar o job de treinamento para o AI Platform Training:
Faça o download do código dos modelos de referência do TensorFlow e acesse o diretório com o código de amostra:
git clone https://github.com/tensorflow/models.git \ --branch=v2.11.0 \ --depth=1 cd models
Crie um arquivo
setup.py
no diretóriomodels
. Isso garante que o comandogcloud ai-platform jobs submit training
inclua todos os subpacotes necessários no diretóriomodels/official
ao criar um tarball do código de treinamento, bem como que o AI Platform Training instale os Conjuntos de dados do TensorFlow como dependências ao executar o job de treinamento. Esse código de treinamento depende dos conjuntos de dados do TensorFlow para carregar os dados do MNIST.Para criar o arquivo
setup.py
, execute o seguinte comando no shell:cat << END > setup.py from setuptools import find_packages from setuptools import setup setup( name='official', install_requires=[ 'tensorflow-datasets~=3.1', 'tensorflow-model-optimization>=0.4.1' ], packages=find_packages() ) END
Envie o job de treinamento usando o comando
gcloud ai-platform jobs submit training
:gcloud ai-platform jobs submit training tpu_mnist_1 \ --staging-bucket=gs://BUCKET_NAME \ --package-path=official \ --module-name=official.vision.image_classification.mnist_main \ --runtime-version=2.11 \ --python-version=3.7 \ --scale-tier=BASIC_TPU \ --region=us-central1 \ -- \ --distribution_strategy=tpu \ --data_dir=gs://tfds-data/datasets \ --model_dir=gs://BUCKET_NAME/tpu_mnist_1_output
Substitua BUCKET_NAME pelo nome de um bucket do Cloud Storage no projeto do Google Cloud. A CLI gcloud envia o código de treinamento empacotado para esse bucket, e o AI Platform Training salva a saída do treinamento no intervalo.
Monitore o job de treinamento. Quando o job for concluído, você poderá ver a saída dele no diretório
gs://BUCKET_NAME/tpu_mnist_1_output
.
Mais detalhes sobre como treinar um modelo na Cloud TPU
Nesta seção, você verá mais informações sobre como configurar um job e treinar um modelo no AI Platform Training com uma Cloud TPU.
Como especificar uma região que ofereça TPUs
Você precisa executar seu job em uma região que tenha TPUs disponíveis. Atualmente, as seguintes regiões dão acesso a TPUs:
us-central1
europe-west4
Para conhecer todas as regiões disponíveis aos serviços do AI Platform Training, incluindo treinamento de modelos e previsão on-line/em lote, leia este guia.
Controle de versões do AI Platform Training e do TensorFlow
As versões de ambiente de execução 1.15, 2.1, 2.2, 2.3, 2.4, 2.5, 2.6, 2.7, 2.8, 2.9 e 2.11 do AI Platform Training estão disponíveis para treinar seus modelos no Cloud TPU. Saiba mais sobre as versões de ambiente de execução do AI Platform Training e as versões correspondentes do TensorFlow.
A política de controle de versões é a mesma da Cloud TPU. Na sua solicitação de job de treinamento, especifique uma versão de ambiente de execução que esteja disponível para TPUs e corresponda à versão do TensorFlow usada no seu código de treinamento.
Como se conectar ao servidor gRPC da TPU
No programa TensorFlow, use TPUClusterResolver
para se conectar ao servidor gRPC da TPU
em execução na VM da TPU.
O guia do TensorFlow para usar TPUs mostra como usar TPUClusterResolver
com
a estratégia de distribuição
TPUStrategy
.
No entanto, é preciso fazer uma alteração importante ao usar TPUClusterResolver
para código executado no AI Platform Training: não forneça nenhum argumento ao
criar a instância TPUClusterResolver
. Quando os argumentos de palavra-chave tpu
, zone
e
project
são definidos com o valor padrão de None
,
o AI Platform Training fornece automaticamente ao resolvedor de cluster os
detalhes de conexão necessários por meio de variáveis de ambiente.
O exemplo a seguir do TensorFlow 2 mostra como inicializar um resolvedor de clusters e uma estratégia de distribuição para treinamento no AI Platform Training:
import tensorflow as tf
resolver = tf.distribute.cluster_resolver.TPUClusterResolver()
tf.config.experimental_connect_to_cluster(resolver)
tf.tpu.experimental.initialize_tpu_system(resolver)
strategy = tf.distribute.experimental.TPUStrategy(resolver)
Como usar TPUs no código do TensorFlow
Para usar as TPUs em uma máquina, use a API TPUStrategy
do TensorFlow 2. O guia do TensorFlow para usar TPUs mostra como fazer isso.
Para treinar com TPUs no TensorFlow 1, use a API TPUEstimator
. O guia do Cloud TPU para a API TPUEstimator
mostra como fazer isso.
A documentação do Cloud TPU também fornece uma lista de operações de nível baixo do TensorFlow disponíveis no Cloud TPU.
Como usar TPUs no código do PyTorch
Para utilizar uma TPU ao aplicar um contêiner de PyTorch
pré-construído,
use o pacote torch_xla
. Saiba como usar torch_xla
para TPU no treinamento na
documentação do PyTorch. Para mais exemplos de
uso de torch_xla
, consulte os tutoriais no repositório GitHub do PyTorch do
GitHub.
Ao treinar usando uma TPU no AI Platform Training, você usa um único dispositivo XLA, e não vários dispositivos XLA.
Consulte também a seção a seguir nesta página sobre como configurar o job de treinamento para PyTorch e TPU.
Como configurar uma máquina personalizada de TPU
Um job de treinamento de TPU é executado em uma configuração de duas VMs. Uma VM (mestre) executa seu código Python. O mestre controla o servidor do TensorFlow em execução em um worker de TPU.
Para usar uma TPU com o AI Platform Training, configure seu job de treinamento para acessar uma máquina ativada para TPU de uma destas três maneiras:
- Use o nível de escalonamento
BASIC_TPU
. Use esse método para acessar os aceleradores da TPU v2. - Use um worker
cloud_tpu
e um tipo de máquina legada para a VM principal. Use esse método para acessar os aceleradores da TPU v2. - Use um worker
cloud_tpu
e um tipo de máquina do Compute Engine para a VM principal. É possível usar esse método para acessar os aceleradores TPU v2 ou TPU v3. Os aceleradores de TPU v3 estão disponíveis na versão Beta.
Máquina ativada para TPU básica
Defina o nível de escalonamento como BASIC_TPU
para conseguir uma VM principal e uma VM de TPU, incluindo uma TPU com oito núcleos TPU v2, como você fez ao executar o exemplo anterior.
Worker de TPU em uma configuração de tipo de máquina legada
Como alternativa, você pode definir uma configuração de máquina personalizada, caso precise de mais recursos de computação na VM mestre:
- Defina o nível de escalonamento como
CUSTOM
. - Configure a VM principal para usar um tipo de máquina legada adequado aos requisitos do job.
- Configure
workerType
comocloud_tpu
para conseguir uma VM de TPU, incluindo um Cloud TPU com oito núcleos TPU v2. - Defina
workerCount
como 1. - Não especifique um servidor de parâmetros ao usar um Cloud TPU. O
serviço rejeitará a solicitação de job se
parameterServerCount
for maior que zero.
No exemplo a seguir, mostramos um arquivo config.yaml
que usa esse
tipo de configuração:
trainingInput:
scaleTier: CUSTOM
masterType: complex_model_m
workerType: cloud_tpu
workerCount: 1
Worker TPU em uma configuração de tipo de máquina do Compute Engine
Também é possível definir uma configuração de máquina personalizada com
um tipo de máquina do Compute
Engine para sua VM mestre
e um acceleratorConfig
anexado à sua VM de TPU.
É possível usar esse tipo de configuração para configurar um worker da TPU com oito núcleos TPU v2 (semelhante a uma configuração sem acceleratorConfig
) ou um worker da TPU com oito núcleos TPU v3 (Beta). Leia mais sobre a diferença entre os
aceleradores TPU v2 e TPU v3.
Usar um tipo de máquina do Compute Engine também oferece mais flexibilidade para configurar sua VM mestre:
- Defina o nível de escalonamento como
CUSTOM
. - Configure a VM mestre para usar um tipo de máquina do Compute Engine que atenda aos seus requisitos de trabalho.
- Defina
workerType
comocloud_tpu
. - Adicione um
workerConfig
com um campoacceleratorConfig
. Dentro deacceleratorConfig
, definatype
comoTPU_V2
ouTPU_V3
ecount
como8
. Não anexe qualquer outro número de núcleos de TPU. - Defina
workerCount
como 1. - Não especifique um servidor de parâmetros ao usar um Cloud TPU. O
serviço rejeitará a solicitação de job se
parameterServerCount
for maior que zero.
No exemplo a seguir, mostramos um arquivo config.yaml
que usa esse
tipo de configuração:
TPU v2
trainingInput:
scaleTier: CUSTOM
masterType: n1-highcpu-16
workerType: cloud_tpu
workerCount: 1
workerConfig:
acceleratorConfig:
type: TPU_V2
count: 8
TPU v3 (Beta)
trainingInput:
scaleTier: CUSTOM
masterType: n1-highcpu-16
workerType: cloud_tpu
workerCount: 1
workerConfig:
acceleratorConfig:
type: TPU_V3
count: 8
Como usar pods de TPU
Um Pod de TPU é um conjunto de dispositivos de TPU conectados por interfaces de rede dedicadas de alta velocidade. Um Pod de TPU pode ter até 2.048 núcleos, permitindo distribuir a carga de processamento entre várias TPUs.
Para usar pods de TPU, primeiro é necessário registrar uma solicitação de aumento de cota.
Os seguintes exemplos de arquivos config.yaml
mostram como usar pods de TPU:
Pods da TPU v2
trainingInput:
scaleTier: CUSTOM
masterType: n1-highcpu-16
workerType: cloud_tpu
workerCount: 1
workerConfig:
acceleratorConfig:
type: TPU_V2_POD
count: 128
Pods da TPU v3
trainingInput:
scaleTier: CUSTOM
masterType: n1-highcpu-16
workerType: cloud_tpu
workerCount: 1
workerConfig:
acceleratorConfig:
type: TPU_V3_POD
count: 32
Há limitações no número de núcleos de pod que podem ser usados para cada tipo de TPU. Configurações disponíveis:
Tipo de pod da TPU | Número de núcleos de pod disponíveis para uso |
---|---|
TPU_V2_POD |
32, 128, 256, 512 |
TPU_V3_POD |
32, 128, 256 |
Para mais detalhes sobre como fazer uso total dos núcleos de TPU, consulte a documentação do Cloud TPU sobre pods de TPU.
Como usar um contêiner do PyTorch criado anteriormente em um worker da TPU
Se você quiser realizar o treinamento PyTorch com uma TPU,
especifique o campo tpuTfVersion
no trainingInput
do job de treinamento.
Defina tpuTfVersion
para corresponder à versão do contêiner PyTorch
pré-criado que você está usando
para o treinamento.
O AI Platform Training é compatível com treinamento de TPUs para os seguintes contêineres pré-criados do PyTorch:
URI da imagem do contêiner | tpuTfVersion |
---|---|
gcr.io/cloud-ml-public/training/pytorch-xla.1-11 |
pytorch-1.11 |
gcr.io/cloud-ml-public/training/pytorch-xla.1-10 |
pytorch-1.10 |
gcr.io/cloud-ml-public/training/pytorch-xla.1-9 |
pytorch-1.9 |
gcr.io/cloud-ml-public/training/pytorch-xla.1-7 |
pytorch-1.7 |
gcr.io/cloud-ml-public/training/pytorch-xla.1-6 |
pytorch-1.6 |
Por exemplo, para treinar usando o contêiner pré-criado PyTorch 1.11, use
o seguinte arquivo config.yaml
para configurar o treinamento:
trainingInput:
scaleTier: CUSTOM
masterType: n1-highcpu-16
masterConfig:
imageUri: gcr.io/cloud-ml-public/training/pytorch-xla.1-11
workerType: cloud_tpu
workerCount: 1
workerConfig:
imageUri: gcr.io/cloud-ml-public/training/pytorch-xla.1-11
tpuTfVersion: pytorch-1.11
acceleratorConfig:
type: TPU_V2
count: 8
Consulte também a seção anterior nesta página sobre Como usar TPUs no código PyTorch.
Como usar um contêiner personalizado em um worker de TPU
Se você quiser executar um contêiner
personalizado no seu worker de TPU, em vez de usar uma das versões
de ambiente de execução do AI Platform Training compatíveis
com TPUs, especifique outro campo de configuração quando enviar seu job de treinamento. Defina
tpuTfVersion
como uma versão do ambiente de
execução que inclua a versão do TensorFlow que o contêiner usa. Especifique uma versão do ambiente de execução compatível para treinamento com TPUs.
Como você está configurando o job para usar um contêiner personalizado, o AI Platform Training não usa o ambiente de execução ao executar o job de treinamento. No entanto, ele precisa que esse campo seja preenchido para preparar adequadamente o worker da TPU para a versão do TensorFlow usada pelo contêiner personalizado.
No exemplo a seguir, mostramos um arquivo config.yaml
com uma configuração
de TPU semelhante àquela da seção anterior, exceto que neste caso a VM
principal e o worker da TPU executam contêineres personalizados diferentes:
TPU v2
trainingInput:
scaleTier: CUSTOM
masterType: n1-highcpu-16
masterConfig:
imageUri: gcr.io/YOUR_PROJECT_ID/your-master-image-name:your-master-tag-name
workerType: cloud_tpu
workerCount: 1
workerConfig:
imageUri: gcr.io/YOUR_PROJECT_ID/your-worker-image-name:your-worker-tag-name
tpuTfVersion: 2.11
acceleratorConfig:
type: TPU_V2
count: 8
TPU v3 (Beta)
trainingInput:
scaleTier: CUSTOM
masterType: n1-highcpu-16
masterConfig:
imageUri: gcr.io/YOUR_PROJECT_ID/your-master-image-name:your-master-tag-name
workerType: cloud_tpu
workerCount: 1
workerConfig:
imageUri: gcr.io/YOUR_PROJECT_ID/your-worker-image-name:your-worker-tag-name
tpuTfVersion: 2.11
acceleratorConfig:
type: TPU_V3
count: 8
Se você usar o
comando gcloud beta ai-platform jobs submit training
para enviar o job de treinamento, poderá especificar o campo da API tpuTfVersion
com a sinalização --tpu-tf-version
, em vez de em um arquivo config.yaml
.
Como usar TPUClusterResolver
após o provisionamento da TPU
Ao usar um contêiner personalizado, você precisa aguardar o provisionamento da TPU
antes de chamar TPUClusterResolver
. O código
de amostra a seguir mostra como lidar com a lógica TPUClusterResolver
:
def wait_for_tpu_cluster_resolver_ready():
"""Waits for `TPUClusterResolver` to be ready and return it.
Returns:
A TPUClusterResolver if there is TPU machine (in TPU_CONFIG). Otherwise,
return None.
Raises:
RuntimeError: if failed to schedule TPU.
"""
tpu_config_env = os.environ.get('TPU_CONFIG')
if not tpu_config_env:
tf.logging.info('Missing TPU_CONFIG, use CPU/GPU for training.')
return None
tpu_node = json.loads(tpu_config_env)
tf.logging.info('Waiting for TPU to be ready: \n%s.', tpu_node)
num_retries = 40
for i in range(num_retries):
try:
tpu_cluster_resolver = (
tf.contrib.cluster_resolver.TPUClusterResolver(
tpu=[tpu_node['tpu_node_name']],
zone=tpu_node['zone'],
project=tpu_node['project'],
job_name='worker'))
tpu_cluster_resolver_dict = tpu_cluster_resolver.cluster_spec().as_dict()
if 'worker' in tpu_cluster_resolver_dict:
tf.logging.info('Found TPU worker: %s', tpu_cluster_resolver_dict)
return tpu_cluster_resolver
except Exception as e:
if i < num_retries - 1:
tf.logging.info('Still waiting for provisioning of TPU VM instance.')
else:
# Preserves the traceback.
raise RuntimeError('Failed to schedule TPU: {}'.format(e))
time.sleep(10)
# Raise error when failed to get TPUClusterResolver after retry.
raise RuntimeError('Failed to schedule TPU.')
Saiba mais sobre o treinamento distribuído com contêineres personalizados.
A seguir
- Saiba mais sobre modelos de treinamento no AI Platform Training.
- Saiba mais sobre como ajustar hiperparâmetros no AI Platform Training, com atenção especial aos detalhes de ajuste com a Cloud TPU.
- Veja outros modelos de referência do Cloud TPU.
- Otimize os modelos para a Cloud TPU seguindo as práticas recomendadas para Cloud TPU.
- Para receber ajuda sobre diagnósticos e correções, consulte a solução de problemas e as perguntas frequentes do Cloud TPU.