介紹
在我之前的文章中,我討論了一種早期的深度學習方法,用於圖像標題生成。如果你有興趣閱讀,可以在這篇文章的結尾找到那篇文章的連結。
今天,我想再次談論圖像標題生成,但這次是使用更先進的神經網絡架構。我將討論的深度學習方法是由劉等人於2021年提出的論文“CPTR: 全變壓器網絡用於圖像標題生成”中所提出的 [1]。具體來說,我將重現論文中提出的模型,並解釋架構背後的理論。不過,請注意,我不會實際演示訓練過程,因為我只想專注於模型架構。
CPTR背後的想法
事實上,CPTR架構的主要想法與早期的圖像標題生成模型完全相同,因為兩者都使用編碼器-解碼器結構。在之前的論文“Show and Tell: A Neural Image Caption Generator” [2] 中,使用的模型分別是GoogLeNet(又名Inception V1)和LSTM作為兩個組件。下圖顯示了Show and Tell論文中提出的模型的示意圖。
儘管具有相同的編碼器-解碼器結構,但CPTR與之前的方法的不同之處在於編碼器和解碼器本身的基礎。在CPTR中,我們將ViT(視覺變壓器)模型的編碼器部分與原始變壓器模型的解碼器部分結合起來。這兩個組件都使用基於變壓器的架構,這正是CPTR名稱的由來:CaPtion TransformeR。
請注意,本文中的討論將與ViT和變壓器高度相關,因此如果你對這兩個主題不熟悉,我強烈建議你閱讀我之前關於這兩個主題的文章。你可以在本文的結尾找到相關連結。
圖2顯示了原始ViT架構的樣子。綠色框內的所有內容都是將被採用作為CPTR編碼器的架構。

接下來,圖3顯示了原始變壓器架構。藍色框內的組件是我們將在CPTR解碼器中實現的層。

如果我們將綠色和藍色框內的組件結合起來,我們將獲得下圖所示的架構。這正是我們將實現的CPTR模型的樣子。這裡的想法是,ViT編碼器(綠色)通過將輸入圖像編碼為特定的張量表示來工作,然後將用作變壓器解碼器(藍色)的基礎,以生成相應的標題。

這就是你目前需要知道的一切。我會在實現過程中進一步解釋更多細節。
模組導入與參數配置
和往常一樣,我們在代碼中需要做的第一件事是導入所需的模組。在這個例子中,我們只導入torch和torch.nn,因為我們即將從頭開始實現模型。
import torch
import torch.nn as nn
接下來,我們將在代碼塊2中初始化一些參數。如果你已經閱讀了我之前關於使用GoogLeNet和LSTM進行圖像標題生成的文章,你會注意到這裡需要初始化的參數多得多。在這篇文章中,我想盡可能地重現CPTR模型,因此將使用論文中提到的參數進行實現。
BATCH_SIZE = 1 #(1)
IMAGE_SIZE = 384 #(2)
IN_CHANNELS = 3 #(3)
SEQ_LENGTH = 30 #(4)
VOCAB_SIZE = 10000 #(5)
EMBED_DIM = 768 #(6)
PATCH_SIZE = 16 #(7)
NUM_PATCHES = (IMAGE_SIZE//PATCH_SIZE) ** 2 #(8)
NUM_ENCODER_BLOCKS = 12 #(9)
NUM_DECODER_BLOCKS = 4 #(10)
NUM_HEADS = 12 #(11)
HIDDEN_DIM = EMBED_DIM * 4 #(12)
DROP_PROB = 0.1 #(13)
我想解釋的第一個參數是BATCH_SIZE,這在標記為#(1)的行中。這個變數分配的數字在我們的情況下並不重要,因為我們實際上不會訓練這個模型。這個參數設置為1,因為PyTorch默認將輸入張量視為一批樣本。這裡我假設我們在一批中只有一個樣本。
接下來,請記住,在圖像標題生成的情況下,我們同時處理圖像和文本。這基本上意味著我們需要為這兩者設置參數。論文中提到模型接受大小為384×384的RGB圖像作為編碼器輸入。因此,我們根據這些信息為IMAGE_SIZE和IN_CHANNELS變數分配值(#(2)和#(3))。另一方面,論文中沒有提到標題的參數。因此,這裡我假設標題的長度不超過30個單詞(#(4)),詞彙大小估計為10000個獨特的單詞(#(5))。
其餘參數與模型配置相關。這裡我們將EMBED_DIM變數設置為768(#(6))。在編碼器端,這個數字表示每個16×16圖像塊的特徵向量的長度(#(7))。同樣的概念也適用於解碼器端,但在那種情況下,特徵向量將表示標題中的單個單詞。更具體地說,關於PATCH_SIZE參數,我們將使用該值來計算輸入圖像中的塊的總數。由於圖像的大小為384×384,因此總共有576個塊(#(8))。
在使用編碼器-解碼器架構時,可以指定要使用的編碼器和解碼器塊的數量。使用更多的塊通常允許模型在準確性方面表現得更好,但相應地,它將需要更多的計算能力。這篇論文的作者決定堆疊12個編碼器塊(#(9))和4個解碼器塊(#(10))。接下來,由於CPTR是一個基於變壓器的模型,因此需要指定編碼器和解碼器內部注意力塊中的注意力頭數量,在這種情況下,作者使用12個注意力頭(#(11))。HIDDEN_DIM參數的值在論文中沒有提到。然而,根據ViT和變壓器的論文,這個參數的配置是EMBED_DIM的4倍(#(12))。論文中也沒有提到丟棄率。因此,我隨意將DROP_PROB設置為0.1(#(13))。
編碼器
在模組和參數設置好之後,現在我們將進入網絡的編碼器部分。在這一部分,我們將逐一實現和解釋圖4中綠色框內的每個組件。
塊嵌入

你可以在上面的圖5中看到,第一步是將輸入圖像分割成塊。這樣做的原因是,與CNN專注於局部模式不同,ViT通過學習這些塊之間的關係來捕捉全局上下文。我們可以使用代碼塊3中顯示的Patcher類來建模這個過程。為了簡化,我在同一類中也包含了塊嵌入塊內的過程。
class Patcher(nn.Module):
def __init__(self):
super().__init__()
#(1)
self.unfold = nn.Unfold(kernel_size=PATCH_SIZE, stride=PATCH_SIZE)
#(2)
self.linear_projection = nn.Linear(in_features=IN_CHANNELS*PATCH_SIZE*PATCH_SIZE,
out_features=EMBED_DIM)
def forward(self, images):
print(f'images\t\t: {images.size()}')
images = self.unfold(images) #(3)
print(f'after unfold\t: {images.size()}')
images = images.permute(0, 2, 1) #(4)
print(f'after permute\t: {images.size()}')
features = self.linear_projection(images) #(5)
print(f'after lin proj\t: {features.size()}')
return features
塊的分割是使用nn.Unfold層完成的(#(1))。在這裡,我們需要將kernel_size和stride參數設置為PATCH_SIZE(16),以便生成的塊不會重疊。此層還會在應用於輸入圖像時自動展平這些塊。同時,nn.Linear層(#(2))用於執行線性投影,即塊嵌入塊所做的過程。通過將out_features參數設置為EMBED_DIM,這一層將每個展平的塊映射到長度為768的特徵向量。
整個過程在你閱讀forward()方法時會更有意義。你可以在同一代碼塊的#(3)行看到,輸入圖像直接由unfold層處理。接下來,我們需要使用permute()方法處理生成的張量(#(4)),以在將其傳遞給linear_projection層之前交換第一和第二個軸(#(5))。此外,我在這裡還打印出每個層之後的張量維度,以便你能更好地理解每一步的變換。
為了檢查我們的Patcher類是否正常工作,我們可以通過網絡傳遞一個虛擬張量。看看下面的代碼塊4,看看我是怎麼做的。
patcher = Patcher()
images = torch.randn(BATCH_SIZE, IN_CHANNELS, IMAGE_SIZE, IMAGE_SIZE)
features = patcher(images)
# 代碼塊4輸出
images : torch.Size([1, 3, 384, 384])
after unfold : torch.Size([1, 768, 576]) #(1)
after permute : torch.Size([1, 576, 768]) #(2)
after lin proj : torch.Size([1, 576, 768]) #(3)
我上面傳遞的張量表示的是一個大小為384×384的RGB圖像。這裡我們可以看到,在展開操作完成後,張量的維度變為1×768×576(#(1)),表示每個576個塊的展平3×16×16塊。不幸的是,這個輸出形狀不符合我們的需求。請記住,在ViT中,我們將圖像塊視為一個序列,因此我們需要交換第一和第二個軸,因為通常,張量的第一維表示時間軸,而第二維表示每個時間步的特徵向量。當執行permute()操作時,我們的張量現在的維度為1×576×768(#(2))。最後,我們將這個張量傳遞給線性投影層,結果張量的形狀保持不變,因為我們將EMBED_DIM參數設置為相同的大小(768)(#(3))。儘管維度相同,但最終張量中包含的信息應該因為線性投影層的可訓練權重的變換而更豐富。
可學習的位置信息嵌入

在輸入圖像成功轉換為一系列塊之後,接下來要做的就是注入所謂的位置信息嵌入張量。這樣做的原因是,沒有位置信息嵌入的變壓器是排列不變的,這意味著它將輸入序列視為其順序無關緊要。有趣的是,由於圖像不是字面上的序列,我們應該將位置信息嵌入設置為可學習的,這樣它就能夠在某種程度上重新排序它認為最適合表示空間信息的塊序列。然而,請記住,這裡的“重新排序”並不意味著我們實際上重新排列序列。相反,它是通過調整嵌入權重來實現的。
實現非常簡單。我們只需要使用nn.Parameter初始化一個張量,其維度設置為與Patcher模型的輸出相匹配,即576×768。此外,別忘了寫上requires_grad=True,以確保該張量是可訓練的。看看下面的代碼塊5,了解詳細信息。
class LearnableEmbedding(nn.Module):
def __init__(self):
super().__init__()
self.learnable_embedding = nn.Parameter(torch.randn(size=(NUM_PATCHES, EMBED_DIM)),
requires_grad=True)
def forward(self):
pos_embed = self.learnable_embedding
print(f'learnable embedding\t: {pos_embed.size()}')
return pos_embed
現在讓我們運行以下代碼塊,以查看我們的LearnableEmbedding類是否正常工作。你可以在打印的輸出中看到,它成功地創建了預期的位置信息嵌入張量。
learnable_embedding = LearnableEmbedding()
pos_embed = learnable_embedding()
# 代碼塊6輸出
learnable embedding : torch.Size([576, 768])
主要編碼器塊

接下來,我們要做的是構建圖7中顯示的主要編碼器塊。這裡你可以看到,這個塊由幾個子組件組成,即自注意力、層正規化、FFN(前饋網絡)和另一個層正規化。下面的代碼塊7a顯示了我如何在EncoderBlock類的__init__()方法中初始化這些層。
class EncoderBlock(nn.Module):
def __init__(self):
super().__init__()
#(1)
self.self_attention = nn.MultiheadAttention(embed_dim=EMBED_DIM,
num_heads=NUM_HEADS,
batch_first=True, #(2)
dropout=DROP_PROB)
self.layer_norm_0 = nn.LayerNorm(EMBED_DIM) #(3)
self.ffn = nn.Sequential( #(4)
nn.Linear(in_features=EMBED_DIM, out_features=HIDDEN_DIM),
nn.GELU(),
nn.Dropout(p=DROP_PROB),
nn.Linear(in_features=HIDDEN_DIM, out_features=EMBED_DIM),
)
self.layer_norm_1 = nn.LayerNorm(EMBED_DIM) #(5)
我之前提到過,ViT的想法是捕捉圖像中塊之間的關係。這個過程是由我在上面的代碼塊中第(1)行初始化的多頭注意力層完成的。這裡需要注意的一點是,我們需要將batch_first參數設置為True(#(2))。這樣做的原因是,注意力層將與我們的張量形狀兼容,其中批次維度(batch_size)位於張量的第0軸。接下來,兩個層正規化層需要分別初始化,如第(3)和第(5)行所示。最後,我們在第(4)行初始化FFN塊,這些層使用nn.Sequential堆疊,遵循以下方程式所定義的結構。

當__init__()方法完成後,我們現在將繼續forward()方法。讓我們看看下面的代碼塊7b。
def forward(self, features): #(1)
residual = features #(2)
print(f'features & residual\t: {residual.size()}')
#(3)
features, self_attn_weights = self.self_attention(query=features,
key=features,
value=features)
print(f'after self attention\t: {features.size()}')
print(f"self attn weights\t: {self_attn_weights.shape}")
features = self.layer_norm_0(features + residual) #(4)
print(f'after norm\t\t: {features.size()}')
residual = features
print(f'\nfeatures & residual\t: {residual.size()}')
features = self.ffn(features) #(5)
print(f'after ffn\t\t: {features.size()}')
features = self.layer_norm_1(features + residual)
print(f'after norm\t\t: {features.size()}')
return features
在這裡,你可以看到輸入張量被命名為features(#(1))。我這樣命名是因為EncoderBlock的輸入是已經通過Patcher和LearnableEmbedding處理的圖像,而不是原始圖像。在執行任何操作之前,請注意,在編碼器塊中有一個分支與主流程分開,然後返回到正規化層。這個分支通常被稱為殘差連接。為了實現這一點,我們需要將原始輸入張量存儲到residual變量中,如我在第(2)行所示。當輸入張量被複製後,我們現在準備使用多頭注意力層處理原始輸入(#(3))。由於這是一個自注意力(而不是交叉注意力),因此這一層的查詢、鍵和值的輸入都來自features張量。接下來,在第(4)行執行層正規化操作,該層的輸入已經包含來自注意力塊的信息以及殘差連接。剩餘的步驟基本上與我剛才解釋的相同,只不過這裡我們用FFN替換了自注意力塊(#(5))。
在接下來的代碼塊中,我將通過傳遞一個大小為1×576×768的虛擬張量來測試EncoderBlock類,模擬來自先前操作的輸出張量。
encoder_block = EncoderBlock()
features = torch.randn(BATCH_SIZE, NUM_PATCHES, EMBED_DIM)
features = encoder_block(features)
以下是整個過程中張量維度的變化。
# 代碼塊8輸出
features & residual : torch.Size([1, 576, 768]) #(1)
after self attention : torch.Size([1, 576, 768])
self attn weights : torch.Size([1, 576, 576]) #(2)
after norm : torch.Size([1, 576, 768])
features & residual : torch.Size([1, 576, 768])
after ffn : torch.Size([1, 576, 768]) #(3)
after norm : torch.Size([1, 576, 768]) #(4)
在這裡,你可以看到最終輸出張量(#(4))的大小與輸入(#(1))相同,這使我們可以堆疊多個編碼器塊,而不必擔心弄亂張量的維度。不僅如此,張量的大小從一開始到最後一層似乎也保持不變。事實上,注意力塊內部實際上執行了很多變換,但我們無法看到,因為整個過程都是由nn.MultiheadAttention層內部完成的。在該層中產生的張量之一是注意力權重(#(2))。這個權重矩陣的大小為576×576,負責存儲有關圖像中每個塊之間關係的信息。此外,張量維度的變化實際上也發生在FFN層內。每個塊的特徵向量最初的長度為768,變為3072,然後立即縮回768(#(3))。然而,這個變換並沒有被打印出來,因為該過程在代碼塊7a的第(4)行用nn.Sequential包裝。
ViT編碼器

在我們完成所有編碼器組件的實現後,現在我們將組裝它們以構建實際的ViT編碼器。我們將在代碼塊9中的Encoder類中完成這個工作。
class Encoder(nn.Module):
def __init__(self):
super().__init__()
self.patcher = Patcher() #(1)
self.learnable_embedding = LearnableEmbedding() #(2)
#(3)
self.encoder_blocks = nn.ModuleList(EncoderBlock() for _ in range(NUM_ENCODER_BLOCKS))
def forward(self, images): #(4)
print(f'images\t\t\t: {images.size()}')
features = self.patcher(images) #(5)
print(f'after patcher\t\t: {features.size()}')
features = features + self.learnable_embedding() #(6)
print(f'after learn embed\t: {features.size()}')
for i, encoder_block in enumerate(self.encoder_blocks):
features = encoder_block(features) #(7)
print(f"after encoder block #{i}\t: {features.shape}")
return features
在__init__()方法中,我們需要做的就是初始化之前創建的所有組件,即Patcher(#(1))、LearnableEmbedding(#(2))和EncoderBlock(#(3))。在這種情況下,EncoderBlock被初始化在nn.ModuleList內,因為我們希望重複NUM_ENCODER_BLOCKS(12)次。在forward()方法中,它最初通過接受原始圖像作為輸入(#(4))來工作。我們然後使用patcher層處理它(#(5)),將圖像分割成小塊並通過線性投影操作轉換它們。然後,通過逐元素相加將可學習的位置信息嵌入張量注入結果輸出(#(6))。最後,我們將其依次傳遞到12個編碼器塊中,使用一個簡單的for循環(#(7))。
現在,在代碼塊10中,我將通過整個編碼器傳遞一個虛擬圖像。請注意,由於我想專注於這個Encoder類的流程,我重新運行了之前創建的類,並將print()函數註釋掉,以便輸出看起來整潔。
encoder = Encoder()
images = torch.randn(BATCH_SIZE, IN_CHANNELS, IMAGE_SIZE, IMAGE_SIZE)
features = encoder(images)
以下是張量的流動情況。在這裡,我們可以看到我們的虛擬輸入圖像成功地通過了網絡中的所有層,包括我們重複12次的編碼器塊。生成的輸出張量現在是上下文感知的,這意味著它已經包含了有關圖像中塊之間關係的信息。因此,這個張量現在準備進一步處理解碼器,這將在後面的部分中討論。
# 代碼塊10輸出
images : torch.Size([1, 3, 384, 384])
after patcher : torch.Size([1, 576, 768])
after learn embed : torch.Size([1, 576, 768])
after encoder block #0 : torch.Size([1, 576, 768])
after encoder block #1 : torch.Size([1, 576, 768])
after encoder block #2 : torch.Size([1, 576, 768])
after encoder block #3 : torch.Size([1, 576, 768])
after encoder block #4 : torch.Size([1, 576, 768])
after encoder block #5 : torch.Size([1, 576, 768])
after encoder block #6 : torch.Size([1, 576, 768])
after encoder block #7 : torch.Size([1, 576, 768])
after encoder block #8 : torch.Size([1, 576, 768])
after encoder block #9 : torch.Size([1, 576, 768])
after encoder block #10 : torch.Size([1, 576, 768])
after encoder block #11 : torch.Size([1, 576, 768])
ViT編碼器(替代方案)
在我們談論解碼器之前,我想先給你們展示一些東西。如果你認為我們上面的做法太複雜,其實你可以使用PyTorch的nn.TransformerEncoderLayer,這樣你就不需要從頭開始實現EncoderBlock類。為此,我將重新實現Encoder類,但這次我將其命名為EncoderTorch。
class EncoderTorch(nn.Module):
def __init__(self):
super().__init__()
self.patcher = Patcher()
self.learnable_embedding = LearnableEmbedding()
#(1)
encoder_block = nn.TransformerEncoderLayer(d_model=EMBED_DIM,
nhead=NUM_HEADS,
dim_feedforward=HIDDEN_DIM,
dropout=DROP_PROB,
batch_first=True)
#(2)
self.encoder_blocks = nn.TransformerEncoder(encoder_layer=encoder_block,
num_layers=NUM_ENCODER_BLOCKS)
def forward(self, images):
print(f'images\t\t\t: {images.size()}')
features = self.patcher(images)
print(f'after patcher\t\t: {features.size()}')
features = features + self.learnable_embedding()
print(f'after learn embed\t: {features.size()}')
features = self.encoder_blocks(features) #(3)
print(f'after encoder blocks\t: {features.size()}')
return features
在上面的代碼塊中,我們基本上做的是,這次我們使用nn.TransformerEncoderLayer(#(1)),這將根據我們傳遞的參數自動創建一個編碼器塊。要重複多次,我們只需使用nn.TransformerEncoder並將數字傳遞給num_layers參數(#(2))。使用這種方法,我們不必像之前那樣在循環中編寫forward傳遞(#(3))。
代碼塊12中的測試代碼與代碼塊10中的代碼完全相同,只不過這裡我使用EncoderTorch類。你也可以看到這裡的輸出基本上與之前的相同。
encoder_torch = EncoderTorch()
images = torch.randn(BATCH_SIZE, IN_CHANNELS, IMAGE_SIZE, IMAGE_SIZE)
features = encoder_torch(images)
# 代碼塊12輸出
images : torch.Size([1, 3, 384, 384])
after patcher : torch.Size([1, 576, 768])
after learn embed : torch.Size([1, 576, 768])
after encoder blocks : torch.Size([1, 576, 768])
解碼器
在我們成功創建了CPTR架構的編碼器部分後,現在我們將談論解碼器。在這一部分,我將實現圖4中藍色框內的每個組件。根據圖示,我們可以看到解碼器接受兩個輸入,即圖像標題的真實值(藍色框下方)和編碼器生成的嵌入塊序列(來自綠色框的箭頭)。重要的是要知道,圖4中繪製的架構旨在說明訓練階段,在該階段,整個標題真實值被輸入到解碼器中。在推理階段,我們只提供一個
正弦位置信息嵌入

如果你查看CPTR模型,你會看到解碼器的第一步是將每個單詞轉換為相應的特徵向量表示,這是通過單詞嵌入塊完成的。然而,由於這一步非常簡單,我們稍後將實現它。現在,讓我們假設這個單詞向量化過程已經完成,因此我們可以進入位置信息嵌入部分。
正如我之前提到的,由於變壓器本質上是排列不變的,我們需要對輸入序列應用位置信息嵌入。與之前的不同,這裡我們使用所謂的正弦位置信息嵌入。我們可以將其視為一種方法,通過分配從正弦波獲得的數字來標記每個單詞向量。這樣做,我們可以期望模型能夠理解單詞的順序,因為波形模式提供的信息。
如果你回顧代碼塊6的輸出,你會看到編碼器中的位置信息嵌入張量的大小為NUM_PATCHES × EMBED_DIM(576×768)。我們基本上想要在解碼器中創建一個大小為SEQ_LENGTH × EMBED_DIM(30×768)的張量,這些值是根據圖11中顯示的方程計算的。這個張量被設置為不可訓練,因為單詞序列必須保持固定的順序以保留其意義。

在這裡,我想快速解釋以下代碼,因為我在之前的變壓器文章中已經更詳細地討論過。一般來說,我們在這裡的基本操作是使用torch.sin()(#(1))和torch.cos()(#(2))創建正弦和餘弦波。然後,使用第(3)和第(4)行的代碼將這兩個張量合併。
class SinusoidalEmbedding(nn.Module):
def forward(self):
pos = torch.arange(SEQ_LENGTH).reshape(SEQ_LENGTH, 1)
print(f"pos\t\t: {pos.shape}")
i = torch.arange(0, EMBED_DIM, 2)
denominator = torch.pow(10000, i/EMBED_DIM)
print(f"denominator\t: {denominator.shape}")
even_pos_embed = torch.sin(pos/denominator) #(1)
odd_pos_embed = torch.cos(pos/denominator) #(2)
print(f"even_pos_embed\t: {even_pos_embed.shape}")
stacked = torch.stack([even_pos_embed, odd_pos_embed], dim=2) #(3)
print(f"stacked\t\t: {stacked.shape}")
pos_embed = torch.flatten(stacked, start_dim=1, end_dim=2) #(4)
print(f"pos_embed\t: {pos_embed.shape}")
return pos_embed
現在我們可以通過運行代碼塊14來檢查上面的SinusoidalEmbedding類是否正常工作。正如之前所預期的那樣,你可以在這裡看到生成的張量的大小為30×768。這個維度與單詞嵌入塊中進行的過程生成的張量相匹配,使它們能夠以逐元素的方式相加。
sinusoidal_embedding = SinusoidalEmbedding()
pos_embed = sinusoidal_embedding()
# 代碼塊14輸出
pos : torch.Size([30, 1])
denominator : torch.Size([384])
even_pos_embed : torch.Size([30, 384])
stacked : torch.Size([30, 384, 2])
pos_embed : torch.Size([30, 768])
前瞻性掩碼

接下來,我要談論解碼器中的被掩碼自注意力層,如上圖所示。我不打算從頭開始編寫注意力機制。相反,我只會實現所謂的前瞻性掩碼,這對自注意力層非常有用,以便在訓練階段它不會關注標題中的後續單詞。
這樣做的方法非常簡單,我們需要做的就是創建一個三角矩陣,其大小設置為與注意力權重矩陣相匹配,即SEQ_LENGTH × SEQ_LENGTH(30×30)。查看create_mask()函數以獲取詳細信息。
def create_mask(seq_length):
mask = torch.tril(torch.ones((seq_length, seq_length))) #(1)
mask[mask == 0] = -float('inf') #(2)
mask[mask == 1] = 0 #(3)
return mask
儘管創建三角矩陣可以簡單地使用torch.tril()和torch.ones()(#(1)),但這裡我們需要進行一些修改,將0值更改為-inf(#(2)),將1更改為0(#(3))。這樣做的原因是,nn.MultiheadAttention層通過逐元素相加應用掩碼。通過將-inf分配給後續單詞,注意力機制將完全忽略它們。同樣,注意力層內部的過程也在我之前的變壓器文章中詳細討論過。
現在,我將運行該函數,seq_length=7,以便你可以看到掩碼的實際樣子。在完整流程中,我們需要將seq_length參數設置為SEQ_LENGTH(30),以便與實際標題長度匹配。
mask_example = create_mask(seq_length=7)
mask_example
# 代碼塊16輸出
tensor([[0., -inf, -inf, -inf, -inf, -inf, -inf],
[0., 0., -inf, -inf, -inf, -inf, -inf],
[0., 0., 0., -inf, -inf, -inf, -inf],
[0., 0., 0., 0., -inf, -inf, -inf],
[0., 0., 0., 0., 0., -inf, -inf],
[0., 0., 0., 0., 0., 0., -inf],
[0., 0., 0., 0., 0., 0., 0.]])
主要解碼器塊

我們可以在上面的圖中看到,解碼器塊的結構比編碼器塊長一些。似乎一切幾乎相同,只是解碼器部分有一個交叉注意力機制和一個額外的層正規化步驟,放置在其後。這個交叉注意力層實際上可以被視為編碼器和解碼器之間的橋樑,因為它用於捕捉標題中每個單詞與輸入圖像中每個塊之間的關係。來自編碼器的兩個箭頭是注意力層的鍵和值輸入,而查詢則來自解碼器本身的前一層。查看下面的代碼塊17a和17b,了解整個解碼器塊的實現。
class DecoderBlock(nn.Module):
def __init__(self):
super().__init__()
#(1)
self.self_attention = nn.MultiheadAttention(embed_dim=EMBED_DIM,
num_heads=NUM_HEADS,
batch_first=True,
dropout=DROP_PROB)
#(2)
self.layer_norm_0 = nn.LayerNorm(EMBED_DIM)
#(3)
self.cross_attention = nn.MultiheadAttention(embed_dim=EMBED_DIM,
num_heads=NUM_HEADS,
batch_first=True,
dropout=DROP_PROB)
#(4)
self.layer_norm_1 = nn.LayerNorm(EMBED_DIM)
#(5)
self.ffn = nn.Sequential(
nn.Linear(in_features=EMBED_DIM, out_features=HIDDEN_DIM),
nn.GELU(),
nn.Dropout(p=DROP_PROB),
nn.Linear(in_features=HIDDEN_DIM, out_features=EMBED_DIM),
)
#(6)
self.layer_norm_2 = nn.LayerNorm(EMBED_DIM)
在__init__()方法中,我們首先使用nn.MultiheadAttention初始化自注意力(#(1))和交叉注意力(#(3))層。這兩個層現在看起來完全相同,但稍後你會在forward()方法中看到它們的區別。三個層正規化操作分別初始化,如第(2)、(4)和(6)行所示,因為每個層將包含不同的正規化參數。最後,FFN層(#(5))與編碼器中的完全相同,基本上遵循圖8中的方程。
談到forward()方法,這裡最初接受三個輸入:features、captions和attn_mask,分別表示來自編碼器的張量、來自解碼器本身的張量和前瞻性掩碼(#(1))。剩餘的步驟與EncoderBlock的步驟有些相似,只不過這裡我們重複了兩次多頭注意力塊。第一個注意力機制將captions作為查詢、鍵和值的參數(#(2))。這樣做是因為我們希望該層捕捉標題張量內部的上下文,因此稱為自注意力。這裡我們還需要將attn_mask參數傳遞給這一層,以便在訓練階段它無法看到後續單詞。第二個注意力機制則不同(#(3))。由於我們希望結合來自編碼器和解碼器的信息,我們需要將captions張量作為查詢,而features張量將作為鍵和值傳遞,因此稱為交叉注意力。在交叉注意力層中不需要前瞻性掩碼,因為在推理階段,模型能夠一次看到整個輸入圖像,而不是逐個查看塊。當張量經過兩個注意力層處理後,我們將其傳遞到前饋網絡中(#(4))。最後,別忘了在每個子組件之後創建殘差連接並應用層正規化步驟。
def forward(self, features, captions, attn_mask): #(1)
print(f"attn_mask\t\t: {attn_mask.shape}")
residual = captions
print(f"captions & residual\t: {captions.shape}")
#(2)
captions, self_attn_weights = self.self_attention(query=captions,
key=captions,
value=captions,
attn_mask=attn_mask)
print(f"after self attention\t: {captions.shape}")
print(f"self attn weights\t: {self_attn_weights.shape}")
captions = self.layer_norm_0(captions + residual)
print(f"after norm\t\t: {captions.shape}")
print(f"\nfeatures\t\t: {features.shape}")
residual = captions
print(f"captions & residual\t: {captions.shape}")
#(3)
captions, cross_attn_weights = self.cross_attention(query=captions,
key=features,
value=features)
print(f"after cross attention\t: {captions.shape}")
print(f"cross attn weights\t: {cross_attn_weights.shape}")
captions = self.layer_norm_1(captions + residual)
print(f"after norm\t\t: {captions.shape}")
residual = captions
print(f"\ncaptions & residual\t: {captions.shape}")
captions = self.ffn(captions) #(4)
print(f"after ffn\t\t: {captions.shape}")
captions = self.layer_norm_2(captions + residual)
print(f"after norm\t\t: {captions.shape}")
return captions
當DecoderBlock類完成後,我們現在可以用以下代碼測試它。
decoder_block = DecoderBlock()
features = torch.randn(BATCH_SIZE, NUM_PATCHES, EMBED_DIM) #(1)
captions = torch.randn(BATCH_SIZE, SEQ_LENGTH, EMBED_DIM) #(2)
look_ahead_mask = create_mask(seq_length=SEQ_LENGTH) #(3)
captions = decoder_block(features, captions, look_ahead_mask)
在這裡,我們假設features是一個包含由編碼器生成的塊嵌入序列的張量(#(1)),而captions是一個嵌入單詞的序列(#(2))。前瞻性掩碼的seq_length參數設置為SEQ_LENGTH(30),以便與標題中的單詞數量相匹配(#(3))。每個步驟後的張量維度顯示在以下輸出中。
# 代碼塊18輸出
attn_mask : torch.Size([30, 30])
captions & residual : torch.Size([1, 30, 768])
after self attention : torch.Size([1, 30, 768])
self attn weights : torch.Size([1, 30, 30]) #(1)
after norm : torch.Size([1, 30, 768])
features : torch.Size([1, 576, 768])
captions & residual : torch.Size([1, 30, 768])
after cross attention : torch.Size([1, 30, 768])
cross attn weights : torch.Size([1, 30, 576]) #(2)
after norm : torch.Size([1, 30, 768])
captions & residual : torch.Size([1, 30, 768])
after ffn : torch.Size([1, 30, 768])
after norm : torch.Size([1, 30, 768])
在這裡,我們可以看到我們的DecoderBlock類正常工作,因為它成功地將輸入張量處理到網絡的最後一層。在這裡,我希望你仔細查看第(1)和第(2)行的注意力權重。根據這兩行,我們可以確認我們的解碼器實現是正確的,因為自注意力層產生的注意力權重的大小為30×30(#(1)),這基本上意味著這一層確實捕捉到了輸入標題內部的上下文。與此同時,交叉注意力層生成的注意力權重矩陣的大小為30×576(#(2)),這表明它成功捕捉了單詞和塊之間的關係。這基本上意味著,在執行交叉注意力操作後,生成的標題張量已經被圖像中的信息豐富化。
變壓器解碼器

現在我們已經成功創建了整個解碼器的所有組件,接下來我要做的就是將它們組合到一個類中。看看下面的代碼塊19a和19b,看看我是怎麼做到的。
class Decoder(nn.Module):
def __init__(self):
super().__init__()
#(1)
self.embedding = nn.Embedding(num_embeddings=VOCAB_SIZE,
embedding_dim=EMBED_DIM)
#(2)
self.sinusoidal_embedding = SinusoidalEmbedding()
#(3)
self.decoder_blocks = nn.ModuleList(DecoderBlock() for _ in range(NUM_DECODER_BLOCKS))
#(4)
self.linear = nn.Linear(in_features=EMBED_DIM,
out_features=VOCAB_SIZE)
如果你將這個Decoder類與代碼塊9中的Encoder類進行比較,你會發現它們在結構上有些相似。在編碼器中,我們使用Patcher將圖像塊轉換為向量,而在解碼器中,我們使用nn.Embedding層將標題中的每個單詞轉換為向量(#(1)),這我之前沒有解釋過。然後,我們初始化位置信息嵌入層,在解碼器中,我們使用正弦而不是可學習的(#(2))。接下來,我們使用nn.ModuleList堆疊多個解碼器塊(#(3))。在第(4)行中寫的線性層在編碼器中不存在,因為它在這裡是必要的,因為它將負責將每個嵌入的單詞映射到長度為VOCAB_SIZE(10000)的向量。稍後,這個向量將包含字典中每個單詞的logit,我們需要做的就是取出包含最高值的索引,即最有可能被預測的單詞。
在forward()方法中的張量流動也與Encoder類中的流動非常相似。在代碼塊19b中,我們將features、captions和attn_mask作為輸入(#(1))。請記住,在這種情況下,captions張量包含原始單詞序列,因此我們需要先用嵌入層將這些單詞向量化(#(2))。接下來,我們使用第(3)行的代碼注入正弦位置信息嵌入張量,然後最終將其依次傳遞到四個解碼器塊中(#(4))。最後,我們將生成的張量傳遞到最後的線性層,以獲得預測的logits(#(5))。
def forward(self, features, captions, attn_mask): #(1)
print(f"features\t\t: {features.shape}")
print(f"captions\t\t: {captions.shape}")
captions = self.embedding(captions) #(2)
print(f"after embedding\t\t: {captions.shape}")
captions = captions + self.sinusoidal_embedding() #(3)
print(f"after sin embed\t\t: {captions.shape}")
for i, decoder_block in enumerate(self.decoder_blocks):
captions = decoder_block(features, captions, attn_mask) #(4)
print(f"after decoder block #{i}\t: {captions.shape}")
captions = self.linear(captions) #(5)
print(f"after linear\t\t: {captions.shape}")
return captions
此時,你可能會想知道為什麼我們不實現softmax激活函數,如圖示中所繪製的。這是因為在訓練階段,softmax通常包含在損失函數中,而在推理階段,最大值的索引將保持不變,無論是否應用softmax。
現在讓我們運行以下測試代碼,以檢查我們的實現是否存在錯誤。之前我提到過,解碼器類的captions輸入是一個原始單詞序列。要模擬這一點,我們可以簡單地創建一個隨機整數序列,範圍在0到VOCAB_SIZE(10000)之間,長度為SEQ_LENGTH(30個單詞)(#(1))。
decoder = Decoder()
features = torch.randn(BATCH_SIZE, NUM_PATCHES, EMBED_DIM)
captions = torch.randint(0, VOCAB_SIZE, (BATCH_SIZE, SEQ_LENGTH)) #(1)
captions = decoder(features, captions, look_ahead_mask)
以下是結果輸出的樣子。在最後一行中,你可以看到線性層生成了一個大小為30×10000的張量,這表明我們的解碼器模型現在能夠預測詞彙中每個單詞在所有30個序列位置的logit分數。
# 代碼塊20輸出
features : torch.Size([1, 576, 768])
captions : torch.Size([1, 30])
after embedding : torch.Size([1, 30, 768])
after sin embed : torch.Size([1, 30, 768])
after decoder block #0 : torch.Size([1, 30, 768])
after decoder block #1 : torch.Size([1, 30, 768])
after decoder block #2 : torch.Size([1, 30, 768])
after decoder block #3 : torch.Size([1, 30, 768])
after linear : torch.Size([1, 30, 10000])
變壓器解碼器(替代方案)
實際上,也可以通過用nn.TransformerDecoderLayer替換DecoderBlock類來簡化代碼,就像我們在ViT編碼器中所做的那樣。下面是如果使用這種方法的代碼。
class DecoderTorch(nn.Module):
def __init__(self):
super().__init__()
self.embedding = nn.Embedding(num_embeddings=VOCAB_SIZE,
embedding_dim=EMBED_DIM)
self.sinusoidal_embedding = SinusoidalEmbedding()
#(1)
decoder_block = nn.TransformerDecoderLayer(d_model=EMBED_DIM,
nhead=NUM_HEADS,
dim_feedforward=HIDDEN_DIM,
dropout=DROP_PROB,
batch_first=True)
#(2)
self.decoder_blocks = nn.TransformerDecoder(decoder_layer=decoder_block,
num_layers=NUM_DECODER_BLOCKS)
self.linear = nn.Linear(in_features=EMBED_DIM,
out_features=VOCAB_SIZE)
def forward(self, features, captions, tgt_mask):
print(f"features\t\t: {features.shape}")
print(f"captions\t\t: {captions.shape}")
captions = self.embedding(captions)
print(f"after embedding\t\t: {captions.shape}")
captions = captions + self.sinusoidal_embedding()
print(f"after sin embed\t\t: {captions.shape}")
#(3)
captions = self.decoder_blocks(tgt=captions,
memory=features,
tgt_mask=tgt_mask)
print(f"after decoder blocks\t: {captions.shape}")
captions = self.linear(captions)
print(f"after linear\t\t: {captions.shape}")
return captions
你會在__init__()方法中看到的主要區別是使用nn.TransformerDecoderLayer和nn.TransformerDecoder,在第(1)和第(2)行,前者用於初始化單個解碼器塊,後者用於重複多次。接下來,forward()方法與Decoder類中的方法大致相似,只不過在解碼器塊上的前向傳播會自動重複四次,而不需要放在循環中(#(3))。你需要注意的是,來自編碼器的張量(features)必須作為memory參數的參數傳遞。而來自解碼器本身的張量(captions)則必須作為tgt參數的輸入。
代碼塊22中的測試代碼基本上與代碼塊20中的代碼相同。在這裡,你可以看到這個模型也生成了大小為30×10000的最終輸出張量。
decoder_torch = DecoderTorch()
features = torch.randn(BATCH_SIZE, NUM_PATCHES, EMBED_DIM)
captions = torch.randint(0, VOCAB_SIZE, (BATCH_SIZE, SEQ_LENGTH))
captions = decoder_torch(features, captions, look_ahead_mask)
# 代碼塊22輸出
features : torch.Size([1, 576, 768])
captions : torch.Size([1, 30])
after embedding : torch.Size([1, 30, 768])
after sin embed : torch.Size([1, 30, 768])
after decoder blocks : torch.Size([1, 30, 768])
after linear : torch.Size([1, 30, 10000])
整個CPTR模型
最後,是時候將我們剛創建的編碼器和解碼器部分放入一個類中,以實際構建CPTR架構。你可以在下面的代碼塊23中看到,實現非常簡單。我們在這裡需要做的就是初始化編碼器(#(1))和解碼器(#(2))組件,然後將原始圖像和相應的標題真實值
新聞來源
本文由 AI 台灣 運用 AI 技術編撰,內容僅供參考,請自行核實相關資訊。
歡迎加入我們的 AI TAIWAN 台灣人工智慧中心 FB 社團,
隨時掌握最新 AI 動態與實用資訊!