Utilizzo delle TPU per l'addestramento del tuo modello

Le TPU (Tensor Processing Unit) sono circuiti integrati per applicazioni specifiche (ASIC) sviluppati 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 fornisce un'interfaccia di gestione dei job che ti evita di dover gestire la TPU autonomamente. In alternativa, puoi utilizzare l'API jobs di AI Platform Training nello stesso modo in cui la utilizzi per l'addestramento su una CPU o una GPU.

Le API TensorFlow di alto livello ti aiutano a eseguire i tuoi modelli sull'hardware Cloud TPU.

Configurazione dell'ambiente Google Cloud

Configura il tuo ambiente Google Cloud seguendo la sezione di configurazione della guida introduttiva.

Autorizzazione di Cloud TPU ad accedere al progetto

Segui questi passaggi per autorizzare il nome dell'account di servizio Cloud TPU associato al tuo progetto Google Cloud:

1. Ottieni il nome dell'account di servizio Cloud TPU chiamando projects.getConfig. Esempio:

   PROJECT_ID=PROJECT_ID
   
   curl -H "Authorization: Bearer $(gcloud auth print-access-token)"  \
       https://ml.googleapis.com/v1/projects/$PROJECT_ID:getConfig
   

2. Salva il valore del campo serviceAccountProject e tpuServiceAccount restituito dall'API.

3. Inizializza l'account di servizio 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
   

Ora aggiungi l'account di servizio Cloud TPU come membro del progetto con il ruolo Agente di servizio Cloud ML. Completa i seguenti passaggi nella console Google Cloud o utilizzando il comando gcloud:

Per ulteriori dettagli sull'assegnazione dei ruoli agli account di servizio, consulta la documentazione di IAM.

Esempio: addestramento di un modello MNIST di esempio

Questa sezione mostra come addestrare un modello MNIST di esempio utilizzando una TPU e la versione 2.11 del runtime. Il job di esempio utilizza il livello di scalabilità BASIC_TPU predefinito per la configurazione della macchina. Le sezioni successive della guida illustrano come configurare una configurazione personalizzata.

Questo esempio presuppone che tu stia utilizzando una shell Bash con l'interfaccia alla gcloud CLI installata. Esegui i seguenti comandi per recuperare il codice e inviare il tuo job di addestramento ad AI Platform Training:

1. Scarica il codice per i modelli di riferimento di TensorFlow e vai alla directory con il codice campione:

   git clone https://github.com/tensorflow/models.git \
     --branch=v2.11.0 \
     --depth=1
   
   cd models
   

2. Crea un file setup.py nella directory models. In questo modo, il comando gcloud ai-platform jobs submit training include tutti i sottopacchetti necessari nella directory models/official quando crea un file tarball del codice di addestramento e garantisce 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 seguente comando 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
   

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

4. Monitora il job di addestramento. Al termine del 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 su AI Platform Training con Cloud TPU.

Specifica una regione che offre TPU

Devi eseguire il job in una regione in cui sono disponibili le TPU. Al momento, le seguenti regioni forniscono accesso alle TPU:

• us-central1
• europe-west4

Per comprendere appieno le regioni disponibili per i servizi AI Platform Training, tra cui l'addestramento dei modelli e le previsioni online/collettive, consulta la guida alle regioni.

Controllo delle versioni di TensorFlow e AI Platform Training

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.

I criteri di gestione delle versioni sono gli stessi di Cloud TPU. Nella richiesta del job di addestramento, assicurati di specificare una versione runtime disponibile per le TPU e che corrisponda alla versione di TensorFlow utilizzata nel codice di addestramento.

Connessione al server gRPC TPU

Nel programma TensorFlow, utilizza TPUClusterResolver per connetterti al server gRPC TPU in esecuzione sulla VM TPU.

La guida di TensorFlow all'utilizzo delle TPU mostra come utilizzare TPUClusterResolver con la strategia di distribuzione TPUStrategy.

Tuttavia, devi apportare un'importante modifica quando utilizzi TPUClusterResolver per il codice eseguito su AI Platform Training: non fornire alcun argomento quando costruisci l'istanza TPUClusterResolver. Quando gli argomenti delle parole chiave tpu, zone e project sono impostati tutti sul valore predefinito None, AI Platform Training fornisce automaticamente al resolver del cluster i dettagli di connessione necessari tramite le variabili di ambiente.

Il seguente esempio di TensorFlow 2 mostra come inizializzare un risolutore di 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. La guida di TensorFlow all'utilizzo delle TPU mostra come eseguire questa operazione.

Per l'addestramento con le TPU in TensorFlow 1, puoi utilizzare l'API TPUEstimator. La guida di Cloud TPU all'API TPUEstimator mostra come eseguire questa operazione.

La documentazione di Cloud TPU fornisce anche un elenco di operazioni di TensorFlow di basso livello disponibili su Cloud TPU.

Utilizzo delle TPU nel codice PyTorch

Per utilizzare una TPU quando utilizzi un contenuto predefinito di PyTorch, utilizza il pacchettotorch_xla. Scopri come utilizzare torch_xla per TPU nell'addestramento nella documentazione di PyTorch. Per altri esempi dell'utilizzo di torch_xla, consulta i tutorial nel repository GitHub di PyTorch XLA

Tieni presente che quando esegui l'addestramento utilizzando 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 TPU viene eseguito su una configurazione di due VM. Una VM (la principale) 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 in modo che acceda a una macchina abilitata per le TPU in uno dei tre modi seguenti:

• Utilizza il livello della bilancia BASIC_TPU. Puoi utilizzare questo metodo per accedere agli acceleratori TPU v2.
• Utilizza un worker cloud_tpu e un tipo di macchina precedente per la VM principale. 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 di base con TPU abilitata

Imposta il livello di scalabilità su BASIC_TPU per ottenere una VM master e una VM TPU, inclusa 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 precedente

In alternativa, puoi configurare una configurazione della macchina personalizzata se hai bisogno di più risorse di calcolo sulla VM principale:

• Imposta il livello della scala su CUSTOM.
• Configura la VM principale in modo che utilizzi un tipo di macchina legacy adatto ai requisiti del tuo job.
• Imposta workerType su cloud_tpu per ottenere una VM TPU che include una Cloud TPU con otto core TPU v2.
• Imposta workerCount su 1.
• Non specificare un server di parametri quando utilizzi una 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

Worker TPU in una configurazione del tipo di macchina Compute Engine

Puoi anche configurare una configurazione della macchina personalizzata con un tipo di macchina Compute Engine per la VM principale e un acceleratorConfig collegato alla VM TPU.

Puoi utilizzare questo tipo di configurazione per impostare 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 principale:

• Imposta il livello della scala su CUSTOM.
• Configura la VM principale in modo che utilizzi un tipo di macchina Compute Engine adatto ai requisiti del tuo job.
• Imposta workerType su cloud_tpu.
• Aggiungi un workerConfig con un campo acceleratorConfig. All'interno di acceleratorConfig, imposta type su TPU_V2 o TPU_V3 e count su 8. Non puoi collegare un altro numero di core TPU.
• Imposta workerCount su 1.
• Non specificare un server di parametri quando utilizzi una 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: n1-highcpu-16
  workerType: cloud_tpu
  workerCount: 1
  workerConfig:
    acceleratorConfig:
      type: TPU_V2
      count: 8

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 di TPU è un insieme di dispositivi TPU connessi da interfacce di rete ad alta velocità dedicate. Un pod di TPU può avere fino a 2048 core TPU, il che ti consente di distribuire il carico di elaborazione su più TPU.

Per utilizzare i pod TPU, devi prima inviare una richiesta di aumento della quota.

I seguenti file config.yaml di esempio mostrano come utilizzare i pod TPU:

trainingInput:
  scaleTier: CUSTOM
  masterType: n1-highcpu-16
  workerType: cloud_tpu
  workerCount: 1
  workerConfig:
    acceleratorConfig:
      type: TPU_V2_POD
      count: 128

trainingInput:
  scaleTier: CUSTOM
  masterType: n1-highcpu-16
  workerType: cloud_tpu
  workerCount: 1
  workerConfig:
    acceleratorConfig:
      type: TPU_V3_POD
      count: 32

Esistono limitazioni al numero di core del pod che possono essere utilizzati per ogni tipo di TPU. Configurazioni disponibili:

Per ulteriori dettagli su come sfruttare al meglio i core del pod di TPU, consulta la documentazione di Cloud TPU sui pod TPU.

Utilizzo di un container PyTorch predefinito su un worker TPU

Se vuoi eseguire l'addestramento di PyTorch con una TPU, devi specificare il campo tpuTfVersion in trainingInput del job di addestramento. Imposta tpuTfVersion in modo che corrisponda alla versione del contenuto predefinito di PyTorch che utilizzi per l'addestramento.

AI Platform Training supporta l'addestramento con TPU per i seguenti container predefiniti di PyTorch:

Ad esempio, per l'addestramento utilizzando il container predefinito PyTorch 1.11, puoi 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 di questa pagina sull'utilizzo delle TPU nel codice PyTorch.

Utilizzo di un container personalizzato su un worker TPU

Se vuoi eseguire un container personalizzato sul tuo worker TPU anziché 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 del runtime che includa la versione di TensorFlow utilizzata dal contenitore. Devi specificare una versione di runtime attualmente supportata per l'addestramento con le TPU.

Poiché configuri il job in modo da 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 per poter preparare correttamente il worker TPU per la versione di TensorFlow utilizzata dal tuo container personalizzato.

L'esempio seguente mostra un file config.yaml con una configurazione TPU simile a quella della sezione precedente, tranne per il fatto che in questo caso la VM principale e il worker TPU eseguono ciascuno contenitori personalizzati diversi:

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

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 contenitore personalizzato, devi attendere il provisioning della TPU prima di poter chiamare TPUClusterResolver per utilizzarla. Il seguente codice campione mostra come gestire la logica di 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 i 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, prestando 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.
• Consulta le domande frequenti e la risoluzione dei problemi di Cloud TPU per ricevere assistenza per la diagnosi e la risoluzione dei problemi.