Efficiente segmentazione delle immagini con PyTorch Parte 4

Efficient Image Segmentation with PyTorch Part 4

Un modello basato su Vision Transformer

In questa serie in 4 parti, implementeremo la segmentazione delle immagini passo dopo passo da zero utilizzando tecniche di deep learning in PyTorch. Questa parte si concentrerà sull’implementazione di un modello basato su Vision Transformer per la segmentazione delle immagini.

Co-autore: Naresh Singh

Sommario dell’articolo

In questo articolo, visiteremo l’architettura Transformer che ha fatto impazzire il mondo del deep learning. Il transformer è un’architettura multimodale che può modellare diverse modalità come il linguaggio, la visione e l’audio.

In questo articolo, impareremo:

• Incontra ChatArena una libreria Python progettata per facilitare la comunicazione e la collaborazione tra più grandi modelli di linguaggio (LLM).
• Cosa è l’allucinazione AI? Cosa va storto con i chatbot AI? Come individuare un’intelligenza artificiale che sta allucinando?
• L’importanza della riproducibilità nell’apprendimento automatico

1. A proposito dell’architettura del transformer e dei concetti chiave coinvolti
2. A comprendere l’architettura del vision transformer
3. A introdurre un modello di vision transformer che viene scritto da zero in modo che si possano apprezzare tutti i blocchi di costruzione e le parti in movimento
4. A seguire un tensore di input alimentato in questo modello e ispezionare come cambia la forma
5. A utilizzare questo modello per eseguire la segmentazione dell’immagine sul dataset Oxford IIIT Pet
6. A osservare i risultati di questo compito di segmentazione
7. A introdurre brevemente il SegFormer, un vision transformer di ultima generazione per la segmentazione semantica

In tutto l’articolo, faremo riferimento a codice e risultati da questo notebook per la formazione del modello. Se desideri riprodurre i risultati, avrai bisogno di una GPU per garantire che il primo notebook completi l’esecuzione in un tempo ragionevole.

Articoli in questa serie

Questa serie è per lettori di tutti i livelli di esperienza con il deep learning. Se vuoi imparare sulla pratica del deep learning e della visione AI, insieme a una solida teoria ed esperienza pratica, sei nel posto giusto! Ci aspettiamo che sia una serie in 4 parti con i seguenti articoli:

1. Concetti e idee
2. Un modello basato su CNN
3. Convoluzioni separate in profondità
4. Un modello basato su Vision Transformer (questo articolo)

Iniziamo il nostro viaggio nei vision transformer con un’introduzione e una comprensione intuitiva dell’architettura del transformer.

L’architettura del Transformer

Possiamo pensare all’architettura del transformer come una composizione di strati intersecanti di comunicazione e calcolo. Questa idea è rappresentata visivamente nella figura 2. Il transformer ha N unità di elaborazione (N è 3 nella figura 2), ognuna delle quali è responsabile dell’elaborazione di una frazione 1/N dell’input. Perché queste unità di elaborazione producano risultati significativi, ciascuna di esse deve avere una visione globale dell’input. Pertanto, il sistema comunica ripetutamente informazioni sui dati in ogni unità di elaborazione a ogni altra unità di elaborazione; mostrato utilizzando le frecce rosse, verdi e blu che vanno da ogni unità di elaborazione a ogni altra unità di elaborazione. Ciò è seguito da alcuni calcoli basati su queste informazioni. Dopo sufficienti ripetizioni di questo processo, il modello è in grado di produrre i risultati desiderati.

Vale la pena notare che la maggior parte delle risorse online discute tipicamente sia l’encoder che il decoder del transformer come presentato nella pubblicazione intitolata “Attention is all you need”. Tuttavia, in questo articolo descriveremo solo la parte dell’encoder del transformer.

Analizziamo più da vicino ciò che costituisce la comunicazione e il calcolo nei transformer.

Comunicazione nei transformer: Attenzione

Nel transformer, la comunicazione è implementata da uno strato noto come strato di attenzione. In PyTorch, questo viene chiamato MultiHeadAttention. Arriveremo alla ragione di quel nome tra poco.

La documentazione dice:

“Permette al modello di partecipare congiuntamente alle informazioni provenienti da diverse sottospazi di rappresentazione, come descritto nell’articolo: Attention is all you need.”

Il meccanismo di attenzione consuma un tensore di input x di forma (Batch, Length, Features) e produce un tensore di forma simile y in modo tale che le caratteristiche per ogni input vengono aggiornate in base a quali altri input nello stesso esempio il tensore sta prestando attenzione. Quindi, le caratteristiche di ogni tensore di lunghezza “Caratteristiche” nell’istanza di dimensione “Lunghezza” vengono aggiornate in base a ogni altro tensore. Ecco dove entra in gioco il costo quadratico del meccanismo di attenzione.

Nel contesto di un transformer di visione, l’input al transformer è un’immagine. Supponiamo che questa sia un’immagine di 128 x 128 (larghezza, altezza). La dividiamo in più piccoli pezzi di dimensioni (16 x 16). Per un’immagine di 128 x 128, otteniamo 64 pezzi (Lunghezza), 8 pezzi in ogni riga e 8 righe di pezzi.

Ognuno di questi 64 pezzi di dimensioni 16 x 16 pixel viene considerato un input separato per il modello transformer. Senza approfondire troppo i dettagli, dovrebbe essere sufficiente pensare a questo processo come guidato da 64 diverse unità di elaborazione, ognuna delle quali sta elaborando un singolo patch di immagine 16×16.

In ogni round, il meccanismo di attenzione in ogni unità di elaborazione è responsabile di guardare il patch di immagine che è responsabile e interrogare ciascuna delle altre 63 unità di elaborazione rimanenti per chiedere loro eventuali informazioni che potrebbero essere rilevanti e utili per aiutarlo a elaborare efficacemente il proprio patch di immagine.

Viene seguito dal calcolo l’operazione di comunicazione tramite attenzione, che vedremo nel prossimo paragrafo.

Calcolo nei transformer: Perceptron Multi-Strato

Il calcolo nei transformer non è altro che un’unità MultiLayerPerceptron (MLP). Questa unità è composta da 2 livelli lineari, con una non linearità GeLU in mezzo. Si può considerare l’utilizzo di altre non linearità. Questa unità proietta per prima l’input a 4 volte la dimensione e lo riproietta a 1x, che è uguale alla dimensione dell’input.

Nel codice che vedremo nel nostro notebook, questa classe viene chiamata MultiLayerPerceptron. Il codice è mostrato di seguito.

class MultiLayerPerceptron(nn.Sequential):    def __init__(self, embed_size, dropout):        super().__init__(            nn.Linear(embed_size, embed_size * 4),            nn.GELU(),            nn.Linear(embed_size * 4, embed_size),            nn.Dropout(p=dropout),        )    # end def# end class

Ora che abbiamo compreso il funzionamento ad alto livello dell’architettura transformer, concentriamo la nostra attenzione sul transformer di visione poiché stiamo per eseguire la segmentazione dell’immagine.

Il Transformer di Visione

Il transformer di visione è stato introdotto per la prima volta dall’articolo intitolato “An Image is Worth 16×16 Words: Transformers for Image Recognition at Scale”. L’articolo discute di come gli autori applichino l’architettura del transformer vanilla al problema della classificazione delle immagini. Ciò viene fatto suddividendo l’immagine in patch di dimensioni 16×16 e trattando ciascuna patch come un token di input per il modello. Il modello di codifica del transformer riceve questi token di input e viene chiesto di prevedere una classe per l’immagine di input.

Nel nostro caso, siamo interessati alla segmentazione dell’immagine. Possiamo considerarlo un compito di classificazione a livello di pixel perché intendiamo prevedere una classe target per pixel.

Apportiamo una piccola ma importante modifica al vanilla vision transformer e sostituiamo l’head MLP per la classificazione con un’head MLP per la classificazione a livello di pixel. Abbiamo un singolo strato lineare in output che è condiviso da ogni patch la cui maschera di segmentazione è prevista dal vision transformer. Questo strato lineare condiviso prevede una maschera di segmentazione per ogni patch che è stata inviata come input al modello.

Nel caso del vision transformer, una patch di dimensioni 16×16 è considerata equivalente a un singolo token di input in un determinato passaggio temporale.

Costruire un’intuizione per le dimensioni dei tensori nei vision transformer

Quando si lavora con deep CNNs, le dimensioni dei tensori che abbiamo utilizzato per la maggior parte erano (N, C H, W), dove le lettere rappresentano le seguenti:

• N: Batch size
• C: Numero di canali
• H: Altezza
• W: Larghezza

Si può notare che questo formato è orientato al processing di immagini 2d, poiché presenta caratteristiche molto specifiche per le immagini.

Con i transformer, d’altra parte, le cose diventano molto più generiche e agnostiche al dominio. Quello che vedremo di seguito si applica alla visione, al testo, al NLP, all’audio o ad altri problemi in cui i dati di input possono essere rappresentati come una sequenza. Vale la pena notare che c’è poca specificità della visione nella rappresentazione dei tensori mentre fluiscono attraverso il nostro vision transformer.

Quando si lavora con i transformer e l’attenzione in generale, ci aspettiamo che i tensori abbiano la seguente forma: (B, T, C), dove le lettere rappresentano le seguenti:

• B: Batch size (stesso delle CNN)
• T: Dimensione temporale o lunghezza della sequenza. Questa dimensione è a volte chiamata L. Nel caso dei vision transformer, ogni patch dell’immagine corrisponde a questa dimensione. Se abbiamo 16 patch dell’immagine, il valore della dimensione T sarà 16
• C: La dimensione del canale o dell’embedding. Questa dimensione è a volte chiamata E. Durante l’elaborazione delle immagini, ogni patch di dimensioni 3x16x16 (Canale, Larghezza, Altezza) viene mappata tramite uno strato di embedding della patch in un embedding di dimensione C. Vedremo come questo viene fatto successivamente.

Entriamo nel dettaglio su come il tensore dell’immagine di input viene mutato e processato lungo la sua strada per prevedere la maschera di segmentazione.

Il percorso di un tensore in un vision transformer

Nelle deep CNNs, il percorso di un tensore appare qualcosa come questo (in una architettura basata su UNet, SegNet o altre CNN).

Il tensore di input è tipicamente di forma (1, 3, 128, 128). Questo tensore passa attraverso una serie di operazioni di convoluzione e max-pooling in cui le sue dimensioni spaziali vengono ridotte e le dimensioni del canale vengono aumentate, di solito di un fattore di 2 ciascuno. Questo è chiamato encoder delle caratteristiche. Dopo di ciò, facciamo l’operazione inversa in cui aumentiamo le dimensioni spaziali e riduciamo le dimensioni del canale. Questo è chiamato decoder delle caratteristiche. Dopo il processo di decodifica, otteniamo un tensore di forma (1, 64, 128, 128). Questo viene proiettato nel numero di canali di output C che desideriamo come (1, C, 128, 128) usando una convoluzione pointwise 1×1 senza bias.

Con i vision transformer, il flusso è molto più complesso. Diamo un’occhiata a un’immagine qui sotto e poi cerchiamo di capire come il tensore trasforma le forme ad ogni passaggio lungo la strada.

Analizziamo ogni passaggio in maggior dettaglio e vediamo come aggiorna la forma del tensore che fluisce attraverso il vision transformer. Per comprendere meglio questo, prendiamo valori concreti per le dimensioni del nostro tensore.

1. Batch Normalization: I tensori di input e output hanno forma (1, 3, 128, 128). La forma rimane invariata, ma i valori sono normalizzati a media zero e varianza unitaria.
2. Immagine a patch: Il tensore di input di forma (1, 3, 128, 128) viene convertito in un patch impilato di immagini 16×16. Il tensore di output ha forma (1, 64, 768).
3. Embedding del patch: Il layer di embedding del patch mappa i 768 canali di input in 512 canali di embedding (per questo esempio). Il tensore di output ha forma (1, 64, 512). Il layer di embedding del patch è essenzialmente solo un layer nn.Linear in PyTorch.
4. Embedding di posizione: Il layer di embedding di posizione non ha un tensore di input, ma contribuisce efficacemente a un parametro apprendibile (tensore addestrabile in PyTorch) della stessa forma dell’embedding del patch. Questo ha forma (1, 64, 512).
5. Addizione: Gli embedding del patch e della posizione vengono aggiunti insieme in modo puntuale per produrre l’input al nostro encoder del vision transformer. Questo tensore ha forma (1, 64, 512). Noterete che il principale lavoro del vision transformer, cioè l’encoder, lascia sostanzialmente questa forma del tensore invariata.
6. Encoder del transformer: Il tensore di input di forma (1, 64, 512) fluisce attraverso più blocchi di encoder del transformer, ognuno dei quali ha più attenzioni (comunicazione) seguite da un layer MLP (calcolo). La forma del tensore rimane invariata come (1, 64, 512).
7. Proiezione lineare di output: Se assumiamo di voler segmentare ogni immagine in 10 classi, allora avremo bisogno che ogni patch di dimensioni 16×16 abbia 10 canali. Il layer nn.Linear per la proiezione di output convertirà ora i 512 canali di embedding in 16x16x10 = 2560 canali di output, e questo tensore avrà aspetto (1, 64, 2560). Nel diagramma sopra C’ = 10. Idealmente, sarebbe un percettrone multi-strato, poiché ” i MLP sono approssimatori di funzioni universali “, ma usiamo un singolo layer lineare poiché questo è un esercizio educativo.
8. Patch all’immagine: Questo layer converte i 64 patch codificati come tensore (1, 64, 2560) in qualcosa che assomiglia a una maschera di segmentazione. Questo può essere 10 immagini a singolo canale, o in questo caso una singola immagine a 10 canali, con ogni canale che è la maschera di segmentazione per una delle 10 classi. Il tensore di output ha forma (1, 10, 128, 128).

Ecco fatto – abbiamo segmentato con successo un’immagine di input usando un vision transformer! Ora, diamo un’occhiata a un esperimento insieme a alcuni risultati.

Vision transformers in azione

Questo notebook contiene tutto il codice per questa sezione.

Per quanto riguarda il codice e la struttura della classe, imita da vicino il diagramma a blocchi sopra. La maggior parte dei concetti menzionati sopra ha una corrispondenza 1:1 con i nomi delle classi in questo notebook.

Ci sono alcuni concetti legati ai layer di attenzione che sono iperparametri critici per il nostro modello. Non abbiamo menzionato nulla sui dettagli dell’attenzione multi-head in precedenza poiché abbiamo detto che è fuori dallo scopo di questo articolo. Consigliamo vivamente di leggere il materiale di riferimento sopra menzionato prima di procedere se non si ha una comprensione di base del meccanismo di attenzione nei transformer.

Abbiamo utilizzato i seguenti parametri del modello per il vision transformer per la segmentazione.

1. 768 dimensioni di embedding per il layer PatchEmbedding
2. 12 blocchi di encoder del Transformer
3. 8 attenzioni multi-head in ogni blocco di encoder del Transformer
4. 20% di dropout nell’attenzione multi-head e nel MLP

Questa configurazione può essere vista nella classe dati Python VisionTransformerArgs.

@dataclassclass VisionTransformerArgs:    """Argomenti per il VisionTransformerForSegmentation."""    image_size: int = 128    patch_size: int = 16    in_channels: int = 3    out_channels: int = 3    embed_size: int = 768    num_blocks: int = 12    num_heads: int = 8    dropout: float = 0.2# fine classe

Una configurazione simile a quella precedente è stata utilizzata durante l’addestramento e la validazione del modello. La configurazione è specificata di seguito.

1. Le tecniche di data augmentation di flip orizzontale casuale e jitter di colore vengono applicate al set di addestramento per prevenire l’overfitting
2. Le immagini vengono ridimensionate a 128×128 pixel in un’operazione di ridimensionamento non conservante l’aspetto
3. Nessuna normalizzazione dell’input viene applicata alle immagini – invece viene utilizzato un livello di normalizzazione batch come primo livello del modello
4. Il modello viene addestrato per 50 epoche utilizzando l’ottimizzatore Adam con un LR di 0,0004 e uno scheduler StepLR che fa decadere il tasso di apprendimento di 0,8x ogni 12 epoche
5. La funzione di perdita della cross-entropia viene utilizzata per classificare un pixel come appartenente a un animale domestico, allo sfondo o a un bordo animale domestico

Il modello ha 86,28 milioni di parametri e ha raggiunto un’accuratezza di validazione del 85,89% dopo 50 epoche di addestramento. Questo è inferiore all’accuratezza del 88,28% ottenuta dal modello CNN profondo dopo 20 epoche di addestramento. Ciò potrebbe essere dovuto a alcuni fattori che devono essere validati sperimentalmente.

1. L’ultimo livello di proiezione di output è un singolo nn.Linear e non un perceptron multi-strato
2. La dimensione della patch 16×16 è troppo grande per catturare dettagli più fini
3. Non abbastanza epoche di addestramento
4. Non abbastanza dati di addestramento – si sa che i modelli transformer hanno bisogno di molti più dati per addestrarsi efficacemente rispetto ai modelli CNN profondi
5. Il tasso di apprendimento è troppo basso

Abbiamo tracciato un gif che mostra come il modello impara a prevedere le maschere di segmentazione per 21 immagini nel set di validazione.

Notiamo qualcosa di interessante nelle prime epoche di addestramento. Le maschere di segmentazione previste hanno alcuni strani artefatti di blocco. L’unica ragione che potremmo immaginare per questo è che stiamo scomponendo l’immagine in patch di dimensioni 16×16 e dopo poche epoche di addestramento, il modello non ha imparato nulla di utile al di là di informazioni molto grossolane riguardanti se questa patch 16×16 è generalmente coperta da un animale domestico o da pixel di sfondo.

Ora che abbiamo visto un vision transformer di base in azione, rivolgiamo la nostra attenzione a un vision transformer all’avanguardia per compiti di segmentazione.

SegFormer: Segmentazione semantica con transformers

L’architettura SegFormer è stata proposta in questo articolo nel 2021. Il transformer che abbiamo visto in precedenza è una versione più semplice dell’architettura SegFormer.

In particolare, il SegFormer:

1. Genera 4 set di immagini con patch di dimensioni 4×4, 8×8, 16×16 e 32×32 invece di una singola immagine con patch di dimensioni 16×16
2. Utilizza 4 blocchi codificatori transformer anziché solo 1. Questo sembra un insieme di modelli
3. Utilizza convoluzioni nelle fasi pre e post di auto-attenzione
4. Non utilizza embedding posizionali
5. Ogni blocco transformer elabora immagini a risoluzione spaziale H/4 x W/4, H/8 x W/8, H/16 x W/16 e H/32, W/32
6. Allo stesso modo, i canali aumentano quando le dimensioni spaziali si riducono. Questo è simile ai CNN profondi
7. Le previsioni a più dimensioni spaziali vengono campionate e quindi unite insieme nel decoder
8. Un MLP combina tutte queste previsioni per fornire una previsione finale
9. La previsione finale è a dimensione spaziale H/4, W/4 e non a H, W.

Conclusione

Nella parte 4 di questa serie, siamo stati introdotti all’architettura transformer e in particolare ai vision transformer. Abbiamo sviluppato una comprensione intuitiva del funzionamento dei vision transformer e del blocco di base coinvolto nelle fasi di comunicazione e calcolo dei vision transformer. Abbiamo visto l’approccio unico basato su patch adottato dai vision transformer per prevedere le maschere di segmentazione e poi combinare insieme le previsioni.

Abbiamo esaminato un esperimento che mostra i vision transformer in azione e siamo stati in grado di confrontare i risultati con gli approcci deep CNN. Sebbene il nostro vision transformer non sia all’avanguardia, è stato in grado di ottenere risultati abbastanza decenti. Abbiamo fornito uno sguardo alle metodologie all’avanguardia come SegFormer.

Dovrebbe essere chiaro a questo punto che i transformer hanno molte più parti in movimento e sono più complessi rispetto agli approcci basati su deep CNN. Dal punto di vista dei FLOPs grezzi, i transformer promettono di essere più efficienti. Nei transformer, l’unica vera matrice che richiede molte risorse computazionali è nn.Linear. Questo viene implementato utilizzando la moltiplicazione matriciale ottimizzata sulla maggior parte delle architetture. Grazie a questa semplicità architettonica, i transformer promettono di essere più facili da ottimizzare e velocizzare rispetto agli approcci basati su deep CNN.

Congratulazioni per essere arrivati fin qui! Siamo contenti che abbiate apprezzato la lettura di questa serie sull’efficace segmentazione delle immagini in PyTorch. Se avete domande o commenti, non esitate a lasciarli nella sezione commenti.

Ulteriori letture

I dettagli del meccanismo di attenzione sono fuori dallo scopo di questo articolo. Inoltre, ci sono numerose risorse di alta qualità a cui potete fare riferimento per comprendere in dettaglio il meccanismo di attenzione. Ecco alcune di quelle che vi consigliamo vivamente.

1. The Illustrated Transformer
2. NanoGPT from scratch using PyTorch

Forniremo i link agli articoli che forniscono maggiori dettagli sui vision transformer di seguito.

1. Implementing Vision Transformer (ViT) in PyTorch: Questo articolo dettaglia l’implementazione di un vision transformer per la classificazione delle immagini in PyTorch. In particolare, la loro implementazione utilizza einops, che evitiamo, poiché si tratta di un esercizio focalizzato sull’educazione (raccomandiamo comunque di imparare e utilizzare einops per la leggibilità del codice). Invece, utilizziamo operatori nativi di PyTorch per permutare e riorganizzare le dimensioni del tensore. Inoltre, ci sono alcuni casi in cui l’autore utilizza Conv2d invece di layer Lineari. Abbiamo voluto costruire un’implementazione dei vision transformer senza l’uso di layer convoluzionali del tutto.
2. Vision Transformer: AI Summer
3. Implementing SegFormer in PyTorch