Le Tensor Processing Unit (TPU) sono ASIC di Google sviluppati su misura da Google e utilizzati per accelerare i carichi di lavoro di machine learning. Puoi eseguire i tuoi job di addestramento su AI Platform Training utilizzando Cloud TPU. AI Platform Training offre un'interfaccia di gestione dei job per aiutarti a non dover gestire autonomamente la TPU.
Puoi invece utilizzare l'API AI Platform Training jobs
nello stesso modo in cui la utilizzi per l'addestramento su una CPU o GPU.
Le API TensorFlow di alto livello ti aiutano a far funzionare i tuoi modelli sull'hardware Cloud TPU.
Configurazione dell'ambiente Google Cloud
Configura l'ambiente Google Cloud seguendo la procedura di configurazione della guida introduttiva.
Autorizzazione di Cloud TPU per accedere al tuo progetto
Follow these steps to authorize the Cloud TPU service account name associated with your Google Cloud project:
Get your Cloud TPU service account name by calling
projects.getConfig
. Example:PROJECT_ID=PROJECT_ID curl -H "Authorization: Bearer $(gcloud auth print-access-token)" \ https://ml.googleapis.com/v1/projects/$PROJECT_ID:getConfig
Save the value of the
serviceAccountProject
andtpuServiceAccount
field returned by the API.Initialize the Cloud TPU service account:
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
Now add the Cloud TPU service account as a member in your project,
with the role Cloud ML Service Agent. Complete the following steps in the
Google Cloud console or using the gcloud
command:
Console
- Log in to the Google Cloud console and choose the project in which you're using the TPU.
- Choose IAM & Admin > IAM.
- Click the Add button to add a member to the project.
- Enter the TPU service account in the Members text box.
- Click the Roles dropdown list.
- Enable the Cloud ML Service Agent role (Service Agents > Cloud ML Service Agent).
gcloud
Set environment variables containing your project ID and the Cloud TPU service account:
PROJECT_ID=PROJECT_ID SVC_ACCOUNT=your-tpu-sa-123@your-tpu-sa.google.com.iam.gserviceaccount.com
Grant the
ml.serviceAgent
role to the Cloud TPU service account:gcloud projects add-iam-policy-binding $PROJECT_ID \ --member serviceAccount:$SVC_ACCOUNT --role roles/ml.serviceAgent
For more details about granting roles to service accounts, see the IAM documentation.
Esempio: addestramento di un modello MNIST di esempio
Questa sezione mostra come addestrare un modello MNIST di esempio utilizzando una TPU e una versione di runtime di 2.11. Il job di esempio utilizza il livello di scalabilità BASIC_TPU
predefinito per la configurazione della macchina.
Le sezioni successive della guida mostrano come impostare una configurazione personalizzata.
Questo esempio presuppone che tu stia utilizzando una shell Bash con gcloud CLI installato. Esegui questi comandi per ottenere il codice e inviare il tuo job di addestramento ad AI Platform Training:
Scarica il codice per i modelli di riferimento di TensorFlow e vai alla directory con il codice di esempio:
git clone https://github.com/tensorflow/models.git \ --branch=v2.11.0 \ --depth=1 cd models
Crea un file
setup.py
nella directorymodels
. Questo garantisce che il comandogcloud ai-platform jobs submit training
includa tutti i sottopacchetti necessari all'interno della directorymodels/official
quando crea un tarball del codice di addestramento e assicura che AI Platform Training installi TensorFlow Datasets come dipendenza quando esegue il job di addestramento. Questo codice di addestramento si basa sui set di dati TensorFlow per caricare i dati MNIST.Per creare il file
setup.py
, esegui il comando seguente nella 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
Invia il job di addestramento utilizzando il 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
Sostituisci BUCKET_NAME con il nome di un bucket Cloud Storage nel progetto Google Cloud. L'interfaccia a riga di comando gcloud carica il codice di addestramento pacchettizzato in questo bucket e AI Platform Training salva l'output dell'addestramento nel bucket.
Monitorare il job di addestramento. Una volta completato il job, puoi visualizzarne l'output nella directory
gs://BUCKET_NAME/tpu_mnist_1_output
.
Scopri di più sull'addestramento di un modello su Cloud TPU
Questa sezione fornisce ulteriori informazioni sulla configurazione di un job e sull'addestramento di un modello in AI Platform Training con Cloud TPU.
Specifica di una regione che offre TPU
Devi eseguire il job in un'area geografica in cui sono disponibili le TPU. Le seguenti regioni attualmente forniscono l'accesso alle TPU:
us-central1
europe-west4
Per comprendere appieno le regioni disponibili per i servizi AI Platform Training, inclusi l'addestramento di modelli e la previsione online/batch, consulta la guida alle regioni.
Controllo delle versioni di TensorFlow e AI Platform
Per l'addestramento dei modelli su Cloud TPU sono disponibili le versioni di runtime di AI Platform Training 1.15, 2.1, 2.2, 2.3, 2.4, 2.5, 2.6, 2.7, 2.8, 2.9 e 2.11. Scopri di più sulle versioni di runtime di AI Platform Training e sulle versioni di TensorFlow corrispondenti.
Il criterio di controllo delle versioni è lo stesso di Cloud TPU. Nella richiesta del job di addestramento, assicurati di specificare una versione di runtime disponibile per le TPU e che corrisponda alla versione di TensorFlow utilizzata nel codice di addestramento.
Connessione al server TPU gRPC
Nel programma TensorFlow, utilizza TPUClusterResolver
per connetterti al server TPU gRPC in esecuzione sulla VM TPU.
La guida all'utilizzo delle TPU in TensorFlow mostra come utilizzare TPUClusterResolver
con la strategia di distribuzione TPUStrategy
.
Tuttavia, devi apportare una modifica importante quando utilizzi TPUClusterResolver
per il codice eseguito su AI Platform Training: non fornire argomenti quando crei l'istanza TPUClusterResolver
. Quando gli argomenti delle parole chiave tpu
, zone
e project
sono impostati sul loro valore predefinito di None
, AI Platform Training fornisce automaticamente al resolver del cluster i dettagli di connessione necessari tramite le variabili di ambiente.
L'esempio TensorFlow 2 seguente mostra come inizializzare un resolver del cluster e una strategia di distribuzione per l'addestramento su 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)
Utilizzo delle TPU nel codice TensorFlow
Per utilizzare le TPU su una macchina, utilizza l'API TPUStrategy
di TensorFlow 2. Nella guida di TensorFlow all'utilizzo delle TPU viene descritto come eseguire questa operazione.
Per l'addestramento con TPU in TensorFlow 1, puoi utilizzare invece l'API
TPUEstimator
. Per informazioni su come eseguire questa operazione, consulta la guida dell'API TPUEstimator
per Cloud TPU.
La documentazione di Cloud TPU fornisce inoltre un elenco delle operazioni di TensorFlow di basso livello disponibili su Cloud TPU.
Utilizzo delle TPU nel codice PyTorch
Per utilizzare una TPU quando utilizzi un contenitore PyTorch predefinito, utilizza il pacchetto torch_xla
. Scopri come utilizzare torch_xla
per l'addestramento TPU nella documentazione di PyTorch. Per altri esempi
di utilizzo di torch_xla
, guarda i tutorial nel repository GitHub
PyTorch XLA
Tieni presente che, quando esegui l'addestramento con una TPU su AI Platform Training, utilizzi un singolo dispositivo XLA, non più dispositivi XLA.
Consulta anche la sezione seguente di questa pagina sulla configurazione del job di addestramento per PyTorch e TPU.
Configurazione di una macchina TPU personalizzata
Un job di addestramento di TPU viene eseguito su una configurazione a due VM. Una VM (il master) esegue il tuo codice Python. Il master gestisce il server TensorFlow in esecuzione su un worker TPU.
Per utilizzare una TPU con AI Platform Training, configura il job di addestramento per l'accesso a una macchina abilitata per TPU in uno dei tre modi seguenti:
- Usa il livello di scalabilità
BASIC_TPU
. Puoi utilizzare questo metodo per accedere agli acceleratori TPU v2. - Utilizza un worker
cloud_tpu
e un tipo di macchina legacy per la VM master. Puoi utilizzare questo metodo per accedere agli acceleratori TPU v2. - Utilizza un worker
cloud_tpu
e un tipo di macchina Compute Engine per la VM master. Puoi utilizzare questo metodo per accedere agli acceleratori TPU v2 o TPU v3. Gli acceleratori TPU v3 sono disponibili in versione beta.
Macchina base compatibile con TPU
Imposta il livello di scalabilità su BASIC_TPU
per ottenere una VM master e una VM TPU che includono una TPU con otto core TPU v2, come hai fatto quando hai eseguito l'esempio precedente.
Worker TPU in una configurazione del tipo di macchina legacy
In alternativa, puoi impostare una configurazione di macchine personalizzate se hai bisogno di più risorse di calcolo sulla VM master:
- Imposta il livello di scalabilità su
CUSTOM
. - Configurare la VM master per utilizzare un tipo di macchina legacy adatta ai tuoi requisiti del job.
- Imposta
workerType
sucloud_tpu
per ottenere una VM TPU che includa una Cloud TPU con otto core TPU v2. - Imposta
workerCount
su 1. - Non specificare un server parametri durante l'utilizzo di Cloud TPU. Il servizio rifiuta la richiesta di job se
parameterServerCount
è maggiore di zero.
L'esempio seguente mostra un file config.yaml
che utilizza questo tipo di configurazione:
trainingInput:
scaleTier: CUSTOM
masterType: complex_model_m
workerType: cloud_tpu
workerCount: 1
Lavoratore TPU in una configurazione del tipo di macchina Compute Engine
Puoi anche impostare una configurazione della macchina personalizzata con un
tipo di macchina
Compute Engine per la VM
principale e una acceleratorConfig
collegata alla
VM TPU.
Puoi utilizzare questo tipo di configurazione per configurare un worker TPU con otto core TPU v2 (simile a una configurazione senza acceleratorConfig
) o un worker TPU con otto core TPU v3 (beta). Scopri di più sulla differenza tra gli acceleratori TPU v2 e TPU v3.
L'utilizzo di un tipo di macchina Compute Engine offre inoltre maggiore flessibilità per la configurazione della VM master:
- Imposta il livello di scalabilità su
CUSTOM
. - Configurare la VM master per utilizzare un tipo di macchina Compute Engine adatto ai tuoi requisiti del job.
- Imposta
workerType
sucloud_tpu
. - Aggiungi un
workerConfig
con un campoacceleratorConfig
. All'interno diacceleratorConfig
, impostatype
suTPU_V2
oTPU_V3
ecount
su8
. Non puoi collegare nessun altro numero di core TPU. - Imposta
workerCount
su 1. - Non specificare un server parametri durante l'utilizzo di Cloud TPU. Il servizio rifiuta la richiesta di job se
parameterServerCount
è maggiore di zero.
L'esempio seguente mostra un file config.yaml
che utilizza questo tipo di configurazione:
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
Utilizzo dei pod TPU
Un pod TPU è una raccolta di dispositivi TPU collegati da interfacce di rete dedicate ad alta velocità. Un pod di TPU può avere fino a 2048 core TPU, in modo da poter distribuire il carico di elaborazione su più TPU.
Per utilizzare i pod TPU, devi prima inviare una richiesta di aumento delle quote.
I seguenti file di esempio config.yaml
mostrano come utilizzare i pod TPU:
Pod TPU v2
trainingInput:
scaleTier: CUSTOM
masterType: n1-highcpu-16
workerType: cloud_tpu
workerCount: 1
workerConfig:
acceleratorConfig:
type: TPU_V2_POD
count: 128
Pod TPU v3
trainingInput:
scaleTier: CUSTOM
masterType: n1-highcpu-16
workerType: cloud_tpu
workerCount: 1
workerConfig:
acceleratorConfig:
type: TPU_V3_POD
count: 32
Esistono limitazioni nel numero di core pod che è possibile utilizzare per ogni tipo di TPU. Configurazioni disponibili:
Tipo di pod TPU | Numero di core pod disponibili |
---|---|
TPU_V2_POD |
32, 128, 256, 512 |
TPU_V3_POD |
32, 128, 256 |
Per maggiori dettagli su come utilizzare al meglio i core dei pod TPU, consulta la documentazione di Cloud TPU sui pod TPU.
Utilizzo di un container PyTorch predefinito su un worker TPU
Se vuoi eseguire un addestramento PyTorch con una TPU, devi specificare il campo tpuTfVersion
nel trainingInput
del job di addestramento.
Imposta tpuTfVersion
in modo che corrisponda alla versione del contenitore PyTorch predefinito che utilizzi per l'addestramento.
AI Platform Training supporta l'addestramento con TPU per i seguenti container PyTorch predefiniti:
URI immagine container | 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 |
Ad esempio, per l'addestramento con il container predefinito PyTorch 1.11, potresti utilizzare il seguente file config.yaml
per configurare l'addestramento:
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
Consulta anche la sezione precedente in questa pagina sull'utilizzo delle TPU nel codice PyTorch.
Utilizzo di un container personalizzato su un worker TPU
Se vuoi eseguire un contenitore personalizzato sul worker TPU invece di utilizzare una delle versioni di runtime di AI Platform Training che supportano le TPU, devi specificare un campo di configurazione aggiuntivo quando invii il job di addestramento. Imposta tpuTfVersion
su una versione runtime che includa la versione di TensorFlow utilizzata dal container. Devi specificare una
versione di runtime attualmente supportata per l'addestramento con
TPU.
Poiché stai configurando il tuo job per utilizzare un container personalizzato, AI Platform Training non utilizza l'ambiente di questa versione di runtime quando esegue il job di addestramento. Tuttavia, AI Platform Training richiede questo campo in modo che possa preparare correttamente il worker di TPU per la versione di TensorFlow utilizzata dal container personalizzato.
L'esempio seguente mostra un file config.yaml
con una configurazione TPU simile a quella della sezione precedente, tranne che in questo caso la VM master e il worker TPU eseguono ciascuno container diversi:
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 utilizzi il comando gcloud beta ai-platform jobs submit training
per inviare il job di addestramento, puoi specificare il campo API tpuTfVersion
con il flag --tpu-tf-version
anziché in un file config.yaml
.
Utilizzo di TPUClusterResolver
dopo il provisioning della TPU
Quando utilizzi un container personalizzato, devi attendere il provisioning della TPU prima di poter chiamare TPUClusterResolver
per utilizzarlo. Il seguente codice di esempio mostra come gestire la logica 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.')
Scopri di più sull'addestramento distribuito con container personalizzati.
Passaggi successivi
- Scopri di più sull'addestramento dei modelli in AI Platform Training.
- Scopri di più sull'ottimizzazione degli iperparametri in AI Platform Training, con particolare attenzione ai dettagli per l'ottimizzazione degli iperparametri con Cloud TPU.
- Esplora altri modelli di riferimento per Cloud TPU.
- Ottimizza i tuoi modelli per Cloud TPU seguendo le best practice per Cloud TPU.
- Per assistenza sulla diagnosi e la risoluzione dei problemi, consulta Risoluzione dei problemi e domande frequenti di Cloud TPU.