Risoluzione dei problemi di TensorFlow - TPU

Questa guida, insieme alle domande frequenti, fornisce assistenza per la risoluzione dei problemi per gli utenti che stanno addestrando modelli TensorFlow su Cloud TPU. Se stai risolvendo problemi di addestramento di Pytorch o JAX, puoi consultare i documenti di risoluzione dei problemi relativi a tali framework:

Per guide più generiche su come utilizzare Cloud TPU, consulta:

Panoramica

I problemi comuni riscontrati con Cloud TPU rientrano nelle seguenti categorie:

  1. Problemi di connessione alla TPU

  2. Debug degli errori comuni

  3. Ridurre l'utilizzo della memoria

  4. Migliorare la velocità di addestramento

  5. Eseguire il debug di cali della precisione del modello

Errore di connessione al server TPU

Questa sezione descrive come risolvere i problemi relativi alle situazioni in cui TensorFlow smette di rispondere o stampa un errore durante la connessione alla TPU. La fase di compilazione del grafico TPU può richiedere molto tempo per i modelli di grandi dimensioni, quindi lascia che lo script venga eseguito per almeno cinque minuti prima di concludere che non ha più risposta.

Il primo passaggio consiste nel verificare se il problema riguarda il server stesso o la pipeline di addestramento di TensorFlow. Per farlo, esegui il tutorial MNIST utilizzando l'URL del tuo server TPU e verifica che funzioni correttamente. Se sono ancora presenti problemi di connessione con il tutorial MNIST, significa che si tratta di un problema relativo al server TPU. In questo caso:

  1. Esegui il comando seguente per elencare le TPU disponibili. Sostituzione di zone e project-id con la tua zona e l'ID progetto.

    (vm)$ gcloud compute tpus list --zone zone --project project-id
    

    Verrà visualizzato un output come:

    NAME       ZONE           ACCELERATOR_TYPE  NETWORK_ENDPOINT   NETWORK  RANGE          STATUS
    demo-tpu   us-central1-b  v2-8              10.240.1.2:8470    default  10.240.1.0  READY

  2. Verifica di trasmettere il valore corretto a --tpu (demo-tpu nell'esempio precedente) e che questa TPU sia elencata come READY.

  3. Se la tua TPU non è elencata come READY o continui a riscontrare problemi di connessione, riavvia manualmente il server con:

    (vm)$ gcloud compute tpus stop $TPU_SERVER_NAME && gcloud compute tpus start $TPU_SERVER_NAME

    Nell'esempio precedente $TPU_SERVER_NAME è demo-tpu. L'operazione potrebbe richiedere diversi minuti.

  4. Esegui di nuovo il comando ... tpus list precedente e attendi che la TPU sia in stato READY. L'operazione potrebbe richiedere diversi minuti.

  5. Prova a eseguire di nuovo il tutorial MNIST.

  6. Se continui a riscontrare problemi con l'esecuzione del tutorial MNIST, richiedi assistenza con uno dei meccanismi descritti in Assistenza.

Se l'esempio MNIST viene eseguito correttamente ma il modello smette di rispondere, è probabile che il problema sia dovuto alla pipeline di addestramento. Per eseguire il debug, inizia sostituendo la TPUStrategy nel codice con la strategia predefinita. Quando utilizzi la strategia predefinita, ovunque tu usi strategy.scope() o strategy.run(), il modello viene eseguito sulla CPU (o sulla GPU, se presente) anziché sulla TPU. Se il modello viene eseguito sulla CPU e non sulla TPU, deve esserci un problema specifico della TPU. Se il problema persiste, ti consigliamo di eseguire il debug del problema sulla CPU.

Perdita della connessione a ssh durante l'addestramento

Si potrebbe verificare un timeout della tua connessione ssh a Cloud TPU durante un addestramento a lunga esecuzione (in particolare se utilizzi Cloud Shell). In quel momento, non esiste alcun output per la console TPU e potrebbe sembrare che la TPU abbia interrotto l'addestramento. Per evitare che questo accada, esegui la sessione di formazione con un multiplexer con terminale o uno strumento di gestione delle sessioni come tmux o screen. La connessione ssh rimarrà attiva a prescindere dalla durata dell'addestramento.

Debug degli errori comuni

Impossibile utilizzare il file system locale

Messaggio di errore

InvalidArgumentError: Unimplemented: File system scheme '[local]' not implemented

Framework e configurazioni interessate

Questo messaggio può essere visualizzato durante l'addestramento con TensorFlow utilizzando l'architettura dei nodi TPU.

Dettagli

Tutti i file di input e la directory del modello devono utilizzare un percorso del bucket di archiviazione cloud (gs://bucket-name/...) e questo bucket deve essere accessibile dal server TPU. Tieni presente che tutta l'elaborazione dei dati e il controllo del modello vengono eseguiti sul server TPU, non sulla macchina locale. Per informazioni su come configurare correttamente Cloud Storage per l'utilizzo con la TPU, consulta la guida alla connessione ai bucket Cloud Storage.

Tipo di dati non supportato

Messaggio di errore

TypeError: DataType is not a supported TPU infeed type.

Framework e configurazioni interessate

Questo messaggio può essere visualizzato durante l'addestramento con TensorFlow utilizzando l'architettura dei nodi TPU.

Dettagli

Al momento, sulla TPU sono supportati solo i tipi di dati tf.float32, tf.int32, tf.bfloat16 e tf.bool. Altri tipi di dati comuni, come tf.uint8, tf.string e tf.int64, devono essere convertiti in uno dei tipi di dati supportati durante la pre-elaborazione dei dati (ossia nella pipeline tf.data.Dataset).

Vedi un esempio di conversione nella funzione decode_image utilizzata nell'addestramento MNIST.

Forme dinamiche non supportate

Messaggio di errore

ValueError: shape [Shape] must have a fixed size for dimension d that is known at graph construction time.

Framework e configurazioni interessate

Questo messaggio viene visualizzato solo durante la compilazione XLA con TensorFlow.

Dettagli

Per eseguire l'esecuzione di un modello sulla TPU, TensorFlow lo compila utilizzando il compilatore XLA. Sebbene questa fase di compilazione migliori in modo significativo la velocità di addestramento e la memoria utilizzata, le forme (dimensioni delle dimensioni) di tutti i tensori del grafico devono essere note al momento della compilazione del grafico. Se al momento della compilazione non è possibile determinare alcuna forma, la compilazione TPU non riesce e restituisce un errore come quello riportato sopra.

Un'operazione comune che restituisce una forma dinamica è dataset.batch(batch_size), poiché il numero di campioni rimanenti in uno stream potrebbe essere inferiore alle dimensioni del batch. Di conseguenza, durante l'addestramento sulla TPU, imposta drop remainder=True per dataset.batch. Questo potrebbe comportare l'eliminazione degli ultimi campioni da un file per garantire che ogni batch abbia una forma statica batch_size. Ad esempio:

dataset = tf.data.Dataset.range(8)
dataset = dataset.batch(3, drop_remainder=True)

Operazioni TensorFlow non disponibili

Messaggio di errore

NotFoundError: No registered 'OpName' OpKernel for XLA_TPU_JIT devices compatible with node

Framework e configurazioni interessate

Questo messaggio può essere visualizzato durante l'addestramento con TensorFlow.

Dettagli

Il modello utilizza un'operazione di TensorFlow non disponibile attualmente sulla TPU.

Per un elenco delle operazioni disponibili sulla TPU, insieme ai piani di assistenza futura e suggerimenti per soluzioni alternative, consulta la guida relativa alle operazioni TensorFlow disponibili.

Messaggio di errore per esaurimento della memoria

Messaggio di errore

ResourceExhaustedError: Ran out of memory in memory space hbm; used: YYY; limit: 7.48G.

Framework e configurazioni interessate

Questo messaggio può essere visualizzato durante l'addestramento con TensorFlow, PyTorch o JAX.

Dettagli

Ogni Cloud TPU è composta da otto core TPU, le TPU v2 hanno 8 GB e le TPU v3 hanno 16 GB di RAM (o HBM, memoria a larghezza di banda elevata). Questa memoria viene utilizzata per archiviare i tensori di peso (variabili) e i tensori di risultati intermedi necessari per il calcolo del gradiente. Se il modello è troppo grande per adattarsi alla RAM TPU, l'inizializzazione non riesce e il messaggio di errore precedente viene stampato. Per ulteriore assistenza, consulta la sezione relativa alla riduzione dell'utilizzo della memoria.

Suggerimenti per ridurre l'utilizzo della memoria:

Problemi che impediscono l'esecuzione

Se TensorFlow rileva un errore durante l'esecuzione di TPU, a volte lo script sembra smettere di rispondere anziché uscire dalla shell. In questo caso, premi CTRL+\ sulla tastiera per attivare SIGQUIT, che causa l'uscita immediata di Python.

Analogamente, la pressione di CTRL+C durante l'esecuzione di TPU non arresta TensorFlow immediatamente, ma attende fino alla fine del ciclo di iterazione corrente per uscire correttamente.

Se si verificano nuovi errori durante la riconnessione alla TPU dopo l'uscita in questo modo, reimposta manualmente il server TPU con il comando:

gcloud compute tpus stop tpu-name --zone=zone && gcloud compute tpus start tpu-name --zone=zone

dove tpu-name viene recuperato dalla prima colonna visualizzata dal comando gcloud compute tpus list e zone è la zona mostrata nella seconda colonna.

Spaziatura interna tensore eccessiva

Possibile causa di problema di memoria

I tensori nella memoria TPU sono imbottiti, ovvero la TPU arrotonda le dimensioni dei tensori archiviati in memoria per eseguire i calcoli in modo più efficiente. Questa spaziatura interna si verifica in modo trasparente a livello di hardware e non influisce sui risultati. Tuttavia, in alcuni casi la spaziatura interna può comportare un aumento significativo dell'utilizzo della memoria e dei tempi di esecuzione.

Come ridurre l'utilizzo della memoria

Il software TPU tenta di disporre i tensori in memoria per massimizzare l'efficienza computazionale e ridurre al minimo la spaziatura interna. Questo processo di layout della memoria è complesso, tuttavia, per ottenere risultati ottimali, il modello deve rispettare la seguente regola generale. Per ridurre al minimo l'overhead di memoria e massimizzare l'efficienza computazionale, deve essere vera una delle seguenti condizioni:

  • Le dimensioni totali del batch devono essere un multiplo di 64 (8 per core TPU) e le dimensioni della caratteristica devono essere un multiplo di 128,

    o

  • Le dimensioni totali del batch devono essere un multiplo di 1024 (128 per core TPU) e le dimensioni della funzionalità devono essere un multiplo di 8.

L'utilizzo di una dimensione batch di 1024 e per le dimensioni delle caratteristiche che sono un multiplo di 128 genera la migliore efficienza, anche se potrebbe non essere possibile per tutti i modelli. Per chiarezza, "dimensione caratteristica" si riferisce alle dimensioni nascoste di un livello completamente collegato o al numero di canali di output in una convoluzione. Non tutti i livelli possono essere conformi a questa regola, in particolare il primo e l'ultimo livello della rete. Non è un problema e si prevede che la maggior parte dei modelli richieda una certa spaziatura interna.

Ridurre l'utilizzo della memoria

Se si verifica un errore di memoria esaurita durante l'esecuzione del modello sulla TPU, devi adottare le misure necessarie per ridurre l'utilizzo della memoria del modello.

I modi più efficaci per ridurre l'utilizzo della memoria sono:

  • Riduci la spaziatura interna eccessiva del tensore
  • Riduci le dimensioni del batch

Dimensioni o modello batch troppo grandi

Possibile causa di problema di memoria

Quando addestri una rete neurale su una CPU, GPU o TPU, l'uso della memoria proviene da due posizioni:

  1. La memoria utilizzata è proporzionale al numero di ponderazioni nel modello.
  2. Memorizzazione delle attivazioni intermedie dal pass in avanti necessarie per calcolare il pass indietro. L'utilizzo della memoria è direttamente proporzionale alle dimensioni del batch, alle dimensioni del livello e al numero di livelli.

Pertanto, la memoria richiesta da un modello dipende in gran parte dalle dimensioni del batch.

La memoria richiesta da un modello dipende dal numero di livelli nella rete.

Il runtime TPU cerca di ottimizzare gli operatori per adattare il modello alla memoria (chiamato rematerializzazione, simile al controllo sfumatura), ma non è sempre in grado di eseguire questa operazione.

Come ridurre l'utilizzo della memoria

Riduci lentamente le dimensioni del batch finché non sono in memoria, assicurandoti che le dimensioni del batch totali siano un multiplo di 64 (la dimensione del batch per core deve essere un multiplo di 8). Tieni presente che le dimensioni batch più grandi sono più efficienti sulla TPU. In genere, un totale di 1024 (128 per core) è un buon punto di partenza.

Se il modello non può essere eseguito sulla TPU anche con una dimensione batch ridotta (ad esempio, 64), prova a ridurre il numero di livelli o le dimensioni del livello.

Migliorare la velocità di addestramento

Se il tuo modello è in grado di essere eseguito correttamente sulla TPU, ma la velocità di addestramento è inferiore al previsto, in questa sezione vengono descritti diversi modi per migliorare la velocità. Consulta la guida alle prestazioni per altri suggerimenti su come migliorare le prestazioni di formazione.

Troppi passaggi per esecuzione per loop di addestramento

Descrizione del problema di rendimento

Il passaggio dell'argomento steps_per_execution a Model.compile determina quanti passaggi di addestramento vengono eseguiti tra i callback dell'host. Ogni callback host richiede una comunicazione significativa tra la CPU host del server TPU e il dispositivo TPU, quindi se steps_per_execution è troppo piccolo, può rallentare l'addestramento.

Come scoprire se il modello è interessato

Se un profilo TPU rivela frequenti callback della CPU host tra i passaggi del dispositivo TPU, l'addestramento può usufruire di un valore steps_per_execution più ampio.

Attenuazione

Imposta un valore maggiore di steps_per_execution. Tieni presente che steps_per_execution può essere impostato su un valore elevato, ma tieni presente che la registrazione dei messaggi e il salvataggio di un checkpoint possono avvenire solo dopo l'esecuzione del numero specificato di passaggi.

Collo di bottiglia dell'elaborazione degli input

Descrizione del problema di rendimento

Durante l'addestramento della TPU su un determinato blocco di dati, la funzione di elaborazione degli input prepara il blocco di dati successivo nella CPU. Se la funzione di input impiega più tempo della funzione del modello, la TPU viene lasciata inattiva mentre la funzione di input recupera i dati.

Come scoprire se il modello è interessato

Per visualizzare l'analisi della pipeline di input in TensorBoard, segui le istruzioni nello strumento di analisi delle pipeline di input di Cloud TPU:

immagine

La pagina di analisi della pipeline di input mostra un breve riepilogo che mostra se il modello è bloccato dall'elaborazione degli input. La stessa pagina mostra inoltre il tempo di esecuzione per operazione, che ti consente di individuare le operazioni problematiche.

Attenuazione

Esistono diverse possibili mitigazioni durante il caricamento dei dati con l'API Dataset:

  1. Archivia i dati come una raccolta di strutture tf.train.Example in file TFRecord e caricali con TFRecordDataset. Per esempi, consulta il tutorial sull'API Dataset o il tutorial su ResNet.
  2. Utilizza dataset.cache() e/o dataset.prefetch() per il buffer dei dati di input. evitando così rallentamenti sporadici nell'accesso ai file creando un collo di bottiglia.
  3. Specifica il parametro num_parallel_calls della funzione dataset.map() per abilitare le operazioni map() con più thread. Un'euristica semplice per il valore di num_parallel_calls consiste nell'utilizzare il numero di core CPU disponibili.
  4. Esegui la pre-elaborazione dei dati offline in modo costoso come costo una tantum, anziché incorrere in costi per ogni periodo di ogni addestramento.

L'intera elaborazione di input viene eseguita sulle CPU che si trovano sul server TPU, non sulla macchina locale, quindi la velocità della macchina locale non è un fattore rilevante.

Tempi di elaborazione lenti e utilizzo MXU ridotto

Descrizione del problema di rendimento

Cloud TPU può eseguire moltiplicazioni di matrici e convoluzioni a velocità incredibilmente elevate. La maggior parte delle altre operazioni TensorFlow ha implementazioni efficienti sulla TPU, ma non sono la forza principale di TPU rispetto ad altri hardware. Pertanto, un modello dovrebbe essere dominato da moltiplicazioni o convoluzioni di matrici per sfruttare appieno la TPU.

Come scoprire se il modello è interessato

I sintomi che vedi in questo caso sono passaggi lenti accompagnati da un basso utilizzo di MXU mostrato quando esegui il profilo del rendimento.

Attenuazione

Cerca di ridurre il numero di operazioni che non sono moltiplicazioni della matrice. Dopo aver ridotto il numero di moltiplicazioni delle mascherine, riesegui il benchmark per verificare se le prestazioni sono accettabili sulle TPU.

Spaziatura interna tensore eccessiva

Descrizione del problema di rendimento

I pad tensori in TPU della memoria sono utilizzati in modo che la TPU possa utilizzare le sue unità di calcolo in modo efficiente. La spaziatura interna può aumentare l'utilizzo della memoria e della larghezza di banda della memoria. Consulta la sezione sulla spaziatura interna del tensore per comprendere e risolvere i problemi di spaziatura interna del tensore.

Velocità effettiva lenta e consumo di memoria ridotto

Descrizione del problema di rendimento

Come regola generale, l'utilizzo di batch di dimensioni maggiori comporta una maggiore velocità di addestramento sulla TPU, in termini di campioni/secondo.

Come scoprire se il modello è interessato

La dimensione del batch di ogni modello deve essere sempre pari ad almeno 64 (8 per core TPU), poiché la TPU rimanda sempre i tensori a questa dimensione. Le dimensioni del batch ideale per l'addestramento sulla TPU sono 1024 (128 per core TPU), poiché questo elimina inefficienze relative a trasferimento e spaziatura interna della memoria.

Attenuazione

La best practice consiste nell'utilizzare la dimensione batch più grande che si adatta alla memoria ed è un multiplo di 64. Il modo più semplice per raggiungere questo risultato è iniziare con il codice 1024 e, se questo causa un errore di memoria insufficiente, prova a ridurre le dimensioni del batch fino a quando il modello non viene eseguito correttamente. La modifica delle dimensioni del batch di un modello può richiedere la regolazione di altri iperparametri per ottenere la stessa precisione del modello, ad esempio la frequenza di apprendimento, ma questo fattore deve essere valutato caso per caso.

Dimensioni del livello troppo piccole

Descrizione del problema di rendimento

Anche quando un modello è dominato da moltiplicazioni o convoluzioni della matrice, la CPU potrebbe non funzionare alla massima efficienza se i tensori di input sono piccoli. Rispetto ad altri hardware, la TPU è più efficiente se le dimensioni del batch e le dimensioni del livello sono maggiori (ad esempio, >= 512).

Come scoprire se il modello è interessato

Come regola generale, le dimensioni dei livelli inferiori a 128 raggiungono una scarsa efficienza sulla TPU, dal momento che 128 è la dimensione nativa dell'unità di moltiplicazione della matrice TPU. Per raggiungere livelli elevati di efficienza, si consiglia di utilizzare una dimensione minima nascosta di 512. Tieni presente che i livelli convoluzionali in genere non devono necessariamente essere grandi quanto quelli completamente collegati per raggiungere lo stesso livello di efficienza.

Attenuazione

Se la motivazione principale per le piccole dimensioni del livello nel modello è la velocità di addestramento, riposiziona i modelli con livelli più grandi sulla TPU. Ad esempio, l'aumento delle dimensioni dell'output di un livello da 256 a 512 può solo aumentare il tempo di addestramento del 20%, anche se il modello esegue il doppio dei calcoli.

Profilazione di modelli a livello di operazione

Spesso è utile misurare i tempi di esecuzione e la memoria utilizzata a livello di operazione per identificare i colli di bottiglia delle prestazioni. Per istruzioni su come eseguire questa operazione,
consulta la guida Strumenti per Cloud TPU: Visualizzatore Trace.

Eseguire il debug di cali della precisione del modello

Uno degli obiettivi dell'ecosistema Cloud TPU è che qualsiasi modello attualmente in fase di addestramento su una CPU o GPU ottiene una precisione molto simile quando viene addestrato sulla TPU, con forse piccole modifiche agli iperparametri come la dimensione del batch e la frequenza di apprendimento. A volte, tuttavia, gli utenti possono osservare un peggioramento della precisione durante l'addestramento dei modelli sulla TPU. Il debug di tali problemi può essere estremamente frustrante a causa della natura casuale dell'addestramento della rete neurale. Questa sezione fornisce indicazioni su come individuare la causa principale di eventuali cali nella precisione del modello quando si porta un modello alla TPU.

Informazioni sul partizionamento orizzontale dei dati (parallelism dei dati)

Uno degli obiettivi principali di TensorFlow è che ogni operazione dovrebbe produrre risultati quasi identici a prescindere che venga eseguita sulla CPU, sulla GPU o sulla TPU. Esistono alcune eccezioni, come le operazioni casuali. In generale, se noti una differenza significativa tra l'output delle operazioni non casuali sulla TPU e sulla CPU, segnalalo come bug.

Tuttavia, per la pipeline di addestramento nel suo insieme, esiste una differenza significativa tra l'addestramento su CPU/GPU e TPU. Quando si esegue l'addestramento su una TPU, TensorFlow esegue il partizionamento orizzontale dei dati,ogni Cloud TPU contiene 8 core TPU che operano come unità di elaborazione indipendenti. Per ogni passaggio dell'addestramento, ogni core TPU riceve un batch di dati, calcola i gradienti di peso, scambia i gradienti con gli altri core TPU, quindi calcola l'aggiornamento del peso. Per impostazione predefinita, la perdita viene calcolata in media tra i core, ma può essere sommata modificando il parametro CrossShardOptimizer.

Se la perdita totale del modello può essere calcolata come media (o somma) di perdite indipendenti per campione, questa procedura è matematicamente equivalente all'addestramento su un singolo batch di grandi dimensioni.

L'operazione più comune che non dipende dal campione indipendente è la normalizzazione batch, che viene eseguita separatamente in ogni batch per core. Ad esempio, se la dimensione del batch totale è 128, la dimensione del batch per core è 16 e ciascuno degli 8 core esegue la normalizzazione batch sui propri 16 campioni. In alcuni casi, la normalizzazione batch su piccoli batch (ad esempio, meno di 32) è stata trovata per screditare la precisione. Nello scenario ideale, le dimensioni totali del batch devono essere grandi, ad esempio da 256 a 1024. Se le dimensioni del batch sono troppo grandi per adattarsi alla memoria, l'effetto del partizionamento orizzontale deve essere valutato caso per caso.

Allenamento determinare

Uno dei motivi per cui è difficile eseguire il debug delle differenze di accuratezza del modello è che tra framework diversi (TensorFlow, PyTorch, JAX), il software di addestramento utilizza diverse inizializzazioni del peso e un ordine casuale dei dati ogni volta che un modello viene addestrato. È utile modificare la procedura di addestramento in modo che sia deterministico, in modo che più esecuzioni generino modelli quasi identici. In questa sezione viene spiegato come eseguire il tutorial MNIST in modo deterministico:

  1. Genera un file di checkpoint iniziale eseguendo un singolo passaggio sulla CPU. Il passaggio viene utilizzato per ottenere l'inizializzazione del peso deterministico. Inoltre, assicurati di utilizzare un seed casuale fisso per qualsiasi funzione casuale nel modello.
# Run training for 1 step to create an initial checkpoint.
python mnist_tpu.py \
  --use_tpu=False \
  --data_dir=${STORAGE_BUCKET}/data/ \
  --model_dir=${STORAGE_BUCKET}/init_output \
  --random_seed=12345 \
  --iterations=1
  --train_steps=1
  1. Modifica le funzioni di ordinamento casuale dei dati nella tua funzione di input per utilizzare un seed casuale. Questa azione è già stata eseguita nel tutorial MNIST. Questo vale per le operazioni di input dei dati perché vengono sempre eseguite sulla CPU. Le operazioni casuali nella funzione modello potrebbero non essere deterministiche tra la TPU e la CPU. Ogni volta che chiami un'operazione casuale, passa un seed fisso per garantire gli stessi risultati tra le esecuzioni. Ad esempio:
# In the flag definitions
tf.flags.DEFINE_integer("batch_size", None, "Random seed for training")

# In the input_fn
if FLAGS.random_seed is not None:
dataset = dataset.shuffle(seed=FLAGS.random_seed)
  1. Esegui lo stesso modello due volte sulla CPU per verificare che l'addestramento sia deterministico. Tieni presente che l'addestramento deve essere eseguito per un numero ragionevole di passaggi (ad esempio 1000), ma non deve essere eseguito per convergere.

    Poiché l'addestramento della CPU viene confrontato con un addestramento di TPU a core singolo, utilizza una dimensione in batch che può essere adattata a un singolo core di TPU (in genere, la dimensione totale in batch divisa per 8). TensorFlow non garantisce il determinismo bit per bit tra le esecuzioni, ma la perdita dovrebbe essere molto vicina:
# Copy the initial weights
gsutil mkdir ${STORAGE_BUCKET}/cpu_output_1
gsutil cp -f ${STORAGE_BUCKET}/init_output/* ${STORAGE_BUCKET}/cpu_output_1
gsutil mkdir ${STORAGE_BUCKET}/cpu_output_2
gsutil cp -f ${STORAGE_BUCKET}/init_output/* ${STORAGE_BUCKET}/cpu_output_2

# Run 1
python mnist_tpu.py \
  --use_tpu=False \
  --data_dir=${STORAGE_BUCKET}/data/ \
  --model_dir=${STORAGE_BUCKET}/cpu_output_1 \
  --batch_size=128 \
  --random_seed=12345 \
  --train_steps=2000 \
  --eval_steps=10

# Output 1
accuracy = 0.9910644, global_step = 1000, loss = 0.025323588

# Run 2
python mnist_tpu.py \
  --use_tpu=False \
  --data_dir=${STORAGE_BUCKET}/data/ \
  --model_dir=${STORAGE_BUCKET}/cpu_output_1 \
  --batch_size=128 \
  --random_seed=12345 \
  --train_steps=2000 \
  --eval_steps=10

# Output 2
accuracy = 0.9910644, global_step = 1000, loss = 0.025323414

Addestramento TPU single-core

Dopo aver eseguito il tutorial MNIST in modo deterministico, il passaggio successivo consiste nel replicare i risultati dell'addestramento CPU sulla TPU, utilizzando un un core TPU per indicare se il problema è correlato al partizionamento orizzontale dei dati o al motore di esecuzione TPU stesso.

Ecco come eseguire addestramento e valutazione single-core sul tutorial MNIST:

# Use the same weight initialization as the CPU
gsutil cp -f ${STORAGE_BUCKET}/init_output/* ${STORAGE_BUCKET}/tpu_output

# Run training for 1000 steps
python mnist.py \
    --use_tpu=True \
    --master=$GRPC_SERVER \
    --train_file=${STORAGE_BUCKET}/data/train.tfrecords \
    --model_dir=${STORAGE_BUCKET}/tpu_output \
    --random_seed=12345 \
    --num_shards=1 \
    --batch_size=128 \
    --train_steps=1000 \
    --eval_steps=10

  accuracy = 0.9910644, global_step = 1000, loss = 0.02514153

La perdita non corrisponderà esattamente al modello con addestramento della CPU, ma dovrebbe essere simile. Se non è simile al modello, potrebbe significare che hai trovato un bug nel motore di esecuzione TPU. Prima di inviare una segnalazione di bug, verifica quanto segue:

  1. Stai passando num_shards=1 a TPUConfig.

  2. Non sono presenti operazioni casuali nella funzione del modello e tutte le operazioni casuali in tale funzione di inserimento vengono testate correttamente.

  3. Stai utilizzando lo stesso file di checkpoint iniziale per l'addestramento di CPU e TPU.

Debug dell'addestramento multi-core della TPU

Se il tuo modello ha la stessa perdita sulla CPU e sulla TPU single-core, probabilmente il problema è uno dei seguenti:

(a) Il deterioramento è dovuto alla naturale variazione casuale durante l'addestramento dei modelli neurali con diverse inizializzazioni.

(b) Il deterioramento è dovuto a un problema relativo al partizionamento orizzontale dei dati sulla TPU.

Per determinare se (a) è il problema, addestra nuovamente il modello completo sulla CPU/GPU e sulla TPU multi-core utilizzando la stessa inizializzazione della ponderazione, come sopra.

Se hai la certezza che il calo dell'accuratezza sia statisticamente significativo, i problemi più probabili relativi al partizionamento orizzontale dei dati sono:

  1. Se il modello utilizza la normalizzazione batch, una dimensione totale batch inferiore a 256 (ad esempio, meno di 32 per core) potrebbe ridurre la precisione.
  2. Le funzioni di perdita in modalità batch sono interessate dal partizionamento orizzontale. Tali funzioni di perdita sono in genere abbastanza specializzate. Ad esempio, Karras et al. 2017 utilizzano un discriminatore batch per l'addestramento di una rete pubblicitaria provocatoria.

Risoluzione dei problemi delle VM TPU

I problemi e le soluzioni che seguono sono applicabili solo alle configurazioni di VM TPU.

risoluzione dei problemi di configurazione di gcloud

Problema

gcloud components update mostra il seguente messaggio di errore:

ERROR: (gcloud.components.update)
You cannot perform this action because the Cloud SDK component manager
is disabled for this installation.
Soluzione

Per utilizzare gcloud con una VM TPU, devi utilizzare un'installazione gcloud che non sia gestita tramite un gestore di pacchetti. Per installare gcloud dal codice sorgente:

sudo apt-get remove google-cloud-sdk
curl -O https://dl.google.com/dl/cloudsdk/channels/rapid/downloads/google-cloud-sdk-311.0.0-linux-x86_64.tar.gz
tar -xzf google-cloud-sdk-311.0.0-linux-x86_64.tar.gz
./google-cloud-sdk/install.sh
source ~/.bashrc
Problema

Quando esegui qualsiasi comando che inizia con gcloud alpha compute tpus tpu-vm, vengono mostrate le seguenti informazioni:

ERROR: (gcloud.alpha.compute.tpus) Invalid choice: 'tpu-vm'.
Soluzione

Questo accade quando il repository dei componenti non è stato aggiornato correttamente. Per verificarlo, esegui gcloud --version. La prima riga dell'output dovrebbe essere "Google Cloud SDK HEAD"; se l'output è diverso, l'aggiornamento non ha avuto luogo. In questo caso, prova ad aggiornare i componenti gcloud con il comando seguente.

gcloud components update

Se continui a riscontrare lo stesso errore, prova a reinstallare gcloud con il comando seguente:

gcloud components reinstall
Problema

Il comando gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} --zone ${ZONE} mostra il seguente messaggio di errore:

Waiting for SSH key to propagate.
ssh: connect to host 34.91.136.59 port 22: Connection timed out
ssh: connect to host 34.91.136.59 port 22: Connection timed out
ssh: connect to host 34.91.136.59 port 22: Connection timed out
ERROR: (gcloud.alpha.compute.tpus.tpu-vm.ssh) Could not SSH into the instance.  It is possible that your SSH key has not propagated to the instance yet. Try running this command again.  If you still cannot connect, verify that the firewall and instance are set to accept ssh traffic.
Soluzione

Potrebbe essersi verificato un problema con la propagazione della chiave SSH. Prova a spostare le chiavi generate automaticamente in una posizione di backup per forzare gcloud a ricrearle:

mv ~/.ssh/google_compute_engine ~/.ssh/old-google_compute_engine
mv ~/.ssh/google_compute_engine.pub ~/.ssh/old-google_compute_engine.pub

Log di debug

I framework Cloud TPU supportati, JAX, PyTorch e TensorFlow accedono alle TPU tramite una libreria condivisa denominata libtpu presente su ogni VM TPU. Questa libreria include il compilatore XLA utilizzato per compilare i programmi TPU, il runtime TPU utilizzato per eseguire i programmi compilati e il driver TPU utilizzato dal runtime per l'accesso di basso livello alla TPU.

La libreria libtpu registra informazioni che possono essere utili per il debug. Per impostazione predefinita, questi log sono scritti in /tmp/tpu_logs su ogni VM Cloud TPU. Puoi impostare le seguenti variabili di ambiente prima di iniziare l'addestramento per modificare il comportamento di logging:

TPU_LOG_DIR: la directory in cui sono scritti i log
Il valore predefinito della posizione della directory è /tmp/tpu_logs. La directory viene creata se non esiste già, ma non vengono create directory principali. Se si verifica un errore durante la ricerca o la creazione della directory specificata, un messaggio viene stampato a stderr, ma il programma non viene interrotto e il logging viene disabilitato. Imposta il nome della directory su "disabled" per disattivare del tutto il logging su disco.
TPU_MIN_LOG_LEVEL: la gravità minima che verrà registrata su disco
Le opzioni sono 0 (INFO), 1 (AVVISO), 2 (ERRORE) e 3 (FATAL). Il valore predefinito è 0.
TPU_STDERR_LOG_LEVEL: la gravità minima che verrà registrata su stderr, oltre al disco, se applicabile
Le opzioni disponibili sono le stesse di TPU_MIN_LOG_LEVEL. Il valore predefinito è 3.
TPU_MAX_LOG_SIZE_MB: la dimensione massima in megabyte di ogni file di log
Un nuovo file di log verrà avviato automaticamente quando quello precedente raggiungerà all'incirca queste dimensioni. Il valore predefinito è 1024.