Elaborazione di immagini di grandi dimensioni con Cloud TPU

Questo documento riguarda l'uso del partizionamento spaziale con Cloud TPU e TPUEstimator per consentire l'addestramento di immagini e video di grandi dimensioni.

La funzionalità di partizionamento spaziale Cloud TPU consente di suddividere una singola immagine tra diversi chip TPU. Il partizionamento spaziale ti consente di scalare facilmente i modelli da eseguire con immagini di input troppo grandi per adattarsi alla memoria disponibile per un singolo core su un chip acceleratore.

Questa funzionalità distribuisce le attivazioni su più core TPU, consentendoti di scalare il tuo modello per utilizzare 2, 4, 8 o anche 16 core per l'addestramento. La scalabilità del numero di processori disponibili per l'addestramento elimina la necessità di ridurre il livello dei dati dell'immagine, offrendo una maggiore precisione e prestazioni più efficienti.

Puoi applicare il partizionamento spaziale ad attività come:

  • Segmentazione dell'immagine della scansione tomografica (CT) del computer 3D
  • analisi dei contenuti video
  • rilevamento di oggetti per la guida autonoma senza limitare i dati delle immagini.

Abilitazione del partizionamento spaziale con TPUEstimator

In Tensorflow, il compilatore di ottimizzazione XLA per Tensorflow gestisce automaticamente le comunicazioni tra tutti i core Cloud TPU. Non è necessario modificare il codice per abilitare il partizionamento spaziale per un modello. TPUEstimator supporta l'API di partizionamento spaziale, tutto ciò che devi fare è configurare il partizionamento di ogni tensore di input in TPUConfig.

immagine

Esempio

Il codice seguente mostra un esempio di TPUConfig di partizionamento spaziale a quattro vie per un modello di classificazione delle immagini. Il tensore è suddiviso in quattro parti lungo la dimensione altezza (supponendo che il tensore abbia la forma [batch, altezza, larghezza, canale]).

tpu_config=tpu_config.TPUConfig(
    iterations_per_loop=100,
    num_cores_per_replica=4,
    per_host_input_for_training=tpu_config.InputPipelineConfig.PER_HOST_V2,
    input_partition_dims=[[1, 4, 1, 1], None]]

Per il partizionamento spaziale, la pipeline di input deve essere in formato tf.data e il valore per_host_input_for_training (train_batch_size) deve essere impostato su PER_HOST_V2. L'elemento num_cores_per_replica specificato determina il numero massimo di partizioni in cui è possibile suddividere un'immagine.

I valori input_partition_dims forniscono un elenco con due elementi: feature_partition_dims e label_partition_dims, che descrivono come partizionare i tensori di input. La struttura di feature_partition_dims e label_partition_dims deve corrispondere alla struttura delle funzionalità e delle etichette di input_fn. Specifica "None" per la partizione dell'etichetta in modo che le etichette non vengano divise.

Esecuzione di modelli di riferimento con partizionamento spaziale

Rilevamento di oggetti 2D

RetinaNet è un modello di rilevamento degli oggetti che individua gli oggetti nelle immagini con un riquadro di delimitazione e classifica gli oggetti identificati. Addestrare il modello con immagini molto grandi è una sfida, dato che l'immagine più grande che può adattarsi a un singolo core Cloud TPU (con un batch per dispositivo 8) ha dimensioni di 1280 x 1280. Puoi addestrare immagini 4 volte più grandi partizionando sistematicamente il modello su 8 core TPU, come mostrato.

Dimensioni dell'immagine Tipo di TPU Core TPU Dimensione del batch globale Tempo di passi
1280x1280 v3-8 8 64 910 ms
2560x2560 v3-64 64 64 822 ms

Segmentazione delle immagini 3D

3D UNet è un popolare modello di segmentazione 3D che è stato ampiamente utilizzato nel dominio dell'immagine medica. La risoluzione originale per le immagini CT può essere di massimo 256 x 256 x 256, che non rientrano in un singolo core TPU, quindi i ricercatori spesso devono sottocampionare le immagini. Con il partizionamento spaziale TPU, puoi adattare direttamente le immagini a risoluzione originale utilizzando il partizionamento spaziale a 16 vie.

Dimensioni dell'immagine Tipo di TPU Core TPU Dimensione del batch globale Tempo di passi
128x128x128 v3-8 8 32 3,428
2560x2560 v3-64 64 32 3,02

Risorse aggiuntive