使用 PyTorch 從頭開始構(gòu)建 CLIP | 對比語言圖像預(yù)訓(xùn)練
在2021年，OpenAI發(fā)布了一篇論文《從自然語言監(jiān)督中學(xué)習(xí)可轉(zhuǎn)移的視覺模型》(https://arxiv.org/pdf/2103.00020)，提出了CLIP(對比語言圖像預(yù)訓(xùn)練)，這是一個強(qiáng)大的深度學(xué)習(xí)模型，旨在以統(tǒng)一的方式理解和解釋圖像和文本。它結(jié)合了視覺和語言編碼器，將文本描述與視覺內(nèi)容聯(lián)系起來。CLIP模型本身不生成圖像的描述，但可以用來評估文本和圖像之間的關(guān)系。例如，你可以提供一張貓的圖片，以及一個標(biāo)簽列表，如“貓”和“狗”，以確定哪個標(biāo)簽與圖片匹配的可能性最高。今天，這篇文章將涵蓋使用PyTorch從頭開始實(shí)現(xiàn)CLIP的過程。
CLIP(對比學(xué)習(xí)-圖像預(yù)訓(xùn)練)
傳統(tǒng)的機(jī)器學(xué)習(xí)模型通常需要大量特定任務(wù)的標(biāo)記數(shù)據(jù)集進(jìn)行微調(diào)。例如，一個訓(xùn)練用來識別狗的模型可能在識別貓方面表現(xiàn)不佳，除非它專門針對貓的圖片進(jìn)行了微調(diào)。
CLIP的架構(gòu)支持零樣本學(xué)習(xí)，這意味著它可以執(zhí)行它沒有直接訓(xùn)練過的任務(wù)，通過利用其在圖像和文本之間學(xué)到的廣泛關(guān)聯(lián)。例如，基于它們的文本描述，它可以對它在訓(xùn)練期間從未見過的圖片進(jìn)行分類。引用他們的論文：“我們在零樣本的情況下匹配原始ResNet-50在ImageNet上的準(zhǔn)確性，而不需要使用它訓(xùn)練時的128萬個訓(xùn)練樣本。”
CLIP有以下我們需要構(gòu)建的組件：
- 文本編碼器
- 圖像編碼器
- 自定義數(shù)據(jù)集(如果你正在訓(xùn)練)
- 對稱損失
文本編碼器
由于我們的主要目標(biāo)是使文本和視覺表示的嵌入對齊，我們將需要一個文本編碼器模型來為圖像的文本描述創(chuàng)建特征。本文不會涵蓋如何從頭開始構(gòu)建文本編碼器，而是直接使用變換器庫來創(chuàng)建編碼器，盡管這將涵蓋CLIP實(shí)現(xiàn)的主要思想。為了簡單起見，使用Distil Bert模型是一個不錯的選擇，因?yàn)樗p量級，性能幾乎和標(biāo)準(zhǔn)BERT模型一樣好，具有類似的基礎(chǔ)架構(gòu)。這是需要記住的一點(diǎn)，我們不是加載預(yù)訓(xùn)練版本。
class TextEncoder(nn.Module):
def __init__(self, embed_dim, proj_dim):
super().__init__()
self.model = DistilBertModel(config=DistilBertConfig())
self.layer_norm = nn.LayerNorm(proj_dim)
def forward(self, input_ids, attention_mask):
x = self.model(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state
return self.layer_norm(x)TextEncoder()類將期望兩個輸入，input_ids和attention_mask，這兩個都將通過分詞器生成。
tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased')
texts = ["This is a sample sentence.", "This is another example."]
inputs= tokenizer(texts, padding=True, truncation=True, return_tensors="pt").to(device)
encoder = ImageEncoder(embed_dim=768, proj_dim=256)
inputs = encoder(inputs['input_ids'], inputs['mask'])現(xiàn)在，TextEncoder的前向傳遞輸出將是(Batch_Size, Token_Size + 1, Embed_Size)，在標(biāo)準(zhǔn)BERT架構(gòu)中，模型的目標(biāo)是兩個任務(wù)，輸出一個額外的CLS_Token，附加到原始令牌前面，通常用于進(jìn)一步微調(diào)分類任務(wù)，以及預(yù)測掩蔽令牌，使用掩蔽令牌前后的所有令牌的信息。
由于我們關(guān)心的是為我們的文本數(shù)據(jù)獲取特征嵌入，我們將只取[CLS]令牌，并將其投影到一個共同的空間，與圖像編碼器的視覺嵌入具有相同的嵌入大小。
class TextEncoder(nn.Module):
def __init__(self, embed_dim, proj_dim):
super().__init__()
self.model = DistilBertModel(config=DistilBertConfig())
self.projection = nn.Linear(embed_dim, proj_dim)
self.layer_norm = nn.LayerNorm(proj_dim)
def forward(self, input_ids, attention_mask):
x = self.model(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state
x = x[:, 0, :] # B, T[cls], E
x = self.projection(x)
return self.layer_norm(x)層歸一化是深度學(xué)習(xí)中非常常見的概念，這不是我第一次解釋它，但讓我們再次解釋一下，我們有一個網(wǎng)絡(luò)的輸入，其中包含來自不同類別或特征的數(shù)據(jù)，因?yàn)樵诿總€訓(xùn)練周期中批次會變化，數(shù)據(jù)的分布也會變化，在一批中分布可能在[0, 2)范圍內(nèi)，而在下一批中它可能有樣本分布在[0, 100]范圍內(nèi)。在訓(xùn)練過程中數(shù)據(jù)分布的變化被稱為協(xié)變量偏移。由于輸入的劇烈變化，輸出也會變化，損失也會變化，如果損失劇烈變化，那么在反向傳播過程中權(quán)重將以更高的幅度更新，導(dǎo)致梯度不平滑。簡而言之，歸一化輸入將限制其在整個訓(xùn)練批次中的分布，因此，損失不會有劇烈變化，將導(dǎo)致更平滑的梯度和更快的訓(xùn)練，幫助模型更多地關(guān)注學(xué)習(xí)特征。
圖像編碼器
CLIP有兩個圖像編碼器選項(xiàng)，ResNet或視覺變換器。我們已經(jīng)開發(fā)了各種視覺變換器，因此將使用標(biāo)準(zhǔn)實(shí)現(xiàn)。如果你想使用ResNet作為圖像編碼器，你可以簡單地用視覺變換器模型替換它，你可以使用PyTorch自己的ResNet模型或timm。
class ImageEncoder(nn.Module):
def __init__(self, base_model, embed_dim, proj_dim):
super().__init__()
self.model = base_model
for param in self.model.parameters():
param.requires_grad = True
self.projection = nn.Linear(embed_dim, proj_dim)
self.layer_norm = nn.LayerNorm(proj_dim)
def forward(self, x):
x = self.projection(self.model(x))
return self.layer_norm(x)上面的編碼器類將圖像張量傳遞給模型，然后將其投影到與文本編碼器輸出相同的共同嵌入空間，后面是一個歸一化層。
自定義數(shù)據(jù)集
現(xiàn)在CLIP是一個(相當(dāng))密集的模型，所以如果你想從頭開始訓(xùn)練它，你必須在一個小數(shù)據(jù)集上訓(xùn)練它。由于本文只涉及如何從頭開始實(shí)現(xiàn)架構(gòu)，我們將不會進(jìn)一步詳細(xì)說明如何創(chuàng)建數(shù)據(jù)集，但為了示例，這可能是你想要做的。
class CustomDataset(Dataset):
def __init__(self, texts, image_paths):
self.image_paths = image_paths
self.texts = texts
tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased')
self.inputs = tokenizer(texts, padding=True, truncation=True, return_tensors="pt")
self.transform = torchvision.transforms.ToTensor()
def __len__(self):
return len(self.texts)
def __getitem__(self, idx):
img_path = self.image_paths[idx]
image = Image.open(img_path)
image = self.transform(image)
caption, mask = self.inputs[idx].items()
return {
"image": image,
"input_ids": caption["input_ids"],
"mask": mask["attention_mask"]
}- image_paths：你選擇的數(shù)據(jù)集中圖像的路徑列表。
- texts：數(shù)據(jù)集中每張圖片的標(biāo)題或文本描述。
自定義數(shù)據(jù)集類在調(diào)用Dataset類時創(chuàng)建分詞器并分詞所有文本，我們使用的是distillbert分詞器，這也是我們的文本編碼器模型。
把所有放在一起
class CLIPModel(nn.Module):
def __init__(self):
super().__init__()
self.device = "cuda" if torch.cuda.is_available() else "cpu"
ViT = VissionTransformer(
num_layers=8,
img_size=224,
emb_size=768,
patch_size=16,
num_head=6,
num_class=768).to(self.device)
self.image_encoder = ImageEncoder(base_model=ViT, embed_dim=768, proj_dim=256)
self.text_encoder = TextEncoder(embed_dim=768, proj_dim=256)ClipModel()類是我們把所有放在一起的地方，架構(gòu)將包括來自圖像和文本編碼器的嵌入，然后用于計(jì)算對稱損失。這是核心實(shí)現(xiàn)的NumPy樣式偽代碼。
在我們的實(shí)現(xiàn)中，我們將在CLIPModel類的前向函數(shù)中計(jì)算損失。第一步是獲取圖像和文本嵌入，然后進(jìn)行交叉乘法以獲得相似性矩陣或logits。回到我們的第二張圖。
logits是通過取圖像和文本嵌入的點(diǎn)積來創(chuàng)建的，由于這篇論文基于對比學(xué)習(xí)，我們的主要目標(biāo)是將文本表示與視覺對齊。那么計(jì)算相似性矩陣有什么幫助呢?
答案是每個從圖像編碼器接收到的圖像令牌(圖5：I_1,I_2,.., I_n; 其中I是嵌入，n是批次大小)乘以文本編碼器接收到的每個令牌。得到最終矩陣(B, Token, Embed)@(B, Embed, Token)→(B, Token, Token)。現(xiàn)在我們的任務(wù)是最大化每個對角線元素(I1T1, I2T2,…, InTn)的值。由于我們想要對齊我們的文本和視覺表示，相應(yīng)的圖像令牌應(yīng)該與其相應(yīng)的文本最高相關(guān)。這就是我們將如何為批次中的所有圖像完成的，但讓我們看看單個令牌。
這里，圖并不是真的不同，我們?nèi)D像嵌入I，并與批次中的每個文本嵌入計(jì)算點(diǎn)積。例如，當(dāng)我們使用I3時，我們希望它與批次中相應(yīng)的文本嵌入T3最強(qiáng)地對齊。理想情況下，I3行中最高的值應(yīng)該是點(diǎn)積I3?T3，以同樣的方式批量處理，看起來就像我們在最大化所有對角線元素，其中每個In與其相應(yīng)的Tn最佳對齊。為了實(shí)現(xiàn)這一點(diǎn)，我們使用一個損失函數(shù)來衡量每一行中最大值與其他值的突出程度。這實(shí)際上是通過取行和列的交叉熵?fù)p失來實(shí)現(xiàn)的。
from vit import VissionTransformer # Importing ViT from previous implementaton (GitHub: Ml-Models)
import numpy as np
import torch.nn.functional as F
class CLIPModel(nn.Module):
def __init__(self):
super().__init__()
self.device = "cuda" if torch.cuda.is_available() else "cpu"
ViT = VissionTransformer(
num_layers=8,
img_size=224,
emb_size=768,
patch_size=16,
num_head=6,
num_class=False).to(self.device)
self.image_encoder = ImageEncoder(base_model=ViT, embed_dim=768, proj_dim=256)
self.text_encoder = TextEncoder(embed_dim=768, proj_dim=256)
self.temperature = nn.Parameter(torch.ones([])*np.log(1/7)).to(self.device)
def forward(self, x):
I_t = self.image_encoder(x["image"])
T_t = self.text_encoder(x["input_ids"], x["mask"])
logits = I_t@T_t.T * torch.exp(self.temperature)
labels = torch.arange(I_t.size(0)).to(self.device)
loss_I = F.cross_entropy(logits.T, labels)
loss_T = F.cross_entropy(logits, labels)
loss = (loss_I + loss_T)/2.0
return loss, logits- 我們得到I_t和T_t(大小：B, Token_Size, Embed_Size)
- 我們通過取點(diǎn)積來計(jì)算logits，如前所述，然后乘以溫度參數(shù)的指數(shù)。如果你熟悉對比學(xué)習(xí)或讀過我關(guān)于DINO(無標(biāo)簽蒸餾)的文章，你可能知道通常使用除以溫度來銳化輸出分布。然而，我們不直接除以溫度，而是乘以一個可訓(xùn)練的張量，該張量使用nn.Parameter()設(shè)置，并初始化為log(1/7)。由于eln(x)=x，那么exp(log(1/T))應(yīng)該是1/T，你可能會想知道我們?yōu)槭裁床缓唵蔚爻艘?/T。原因是使用log(1/T)可以讓優(yōu)化器在訓(xùn)練期間更容易計(jì)算和更新梯度。這種方法在深度學(xué)習(xí)中是一種常見的做法，因?yàn)樗梢詭砀交挠?xùn)練和更穩(wěn)定的模型權(quán)重更新。
- 標(biāo)簽簡單地用批次大小生成([0, 1,..N])。正如我們之前討論的，目標(biāo)是最大化每個對角線元素(i1T1, i2T2,..inTn)，因此整個矩陣中每一行的標(biāo)簽是[0, 1, 2, ..N]，對應(yīng)于哪一行的元素應(yīng)該是最大的。
- 如偽代碼所述，嵌入已經(jīng)歸一化，但我們不需要這樣做，因?yàn)槲覀冊诜祷貓D像和文本編碼器的輸出時已經(jīng)應(yīng)用了層歸一化。
- 按照偽代碼，計(jì)算行和列的交叉熵?fù)p失。我們通過傳遞logits的轉(zhuǎn)置和正常的logits以及標(biāo)簽，取兩個損失的平均值，現(xiàn)在我們有最終的損失結(jié)果。
設(shè)置模型
texts = ["This is a sample sentence.", "This is another example."]
# You can Use a CustomDataset as we Implemented above for training
train_data = CustomDataset(texts, image_path)
train_loader = DataLoader(train_data, batch_size, shuffle=True)
# Example Usage
device = 'cuda' if torch.cuda.is_available() else 'cpu'
tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased')
inputs= tokenizer(texts, padding=True, truncation=True, return_tensors="pt").to(device)
test = {
"image" : torch.rand(2, 3, 224, 224).to(device),
"input_ids" : inputs["input_ids"],
"mask" : inputs["attention_mask"]
}
model = CLIPModel().to(device)
loss, logits = model(test)
print("Loss:", loss, "Logits:", logits)源碼鏈接：https://github.com/mishra-18/ML-Models/blob/main/clip.py
