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.
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
- Guida al partizionamento spaziale: ulteriori istruzioni su come configurare il partizionamento spaziale.
- Utilizzo dell'API TPUEstimator su Cloud TPU