使用 PyTorch 從頭開始構建 CLIP | 對比語言圖像預訓練

作者：二旺 2024-09-26 16:51:23

CLIP模型本身不生成圖像的描述，但可以用來評估文本和圖像之間的關系。今天，這篇文章將涵蓋使用PyTorch從頭開始實現CLIP的過程。

在2021年，OpenAI發布了一篇論文《從自然語言監督中學習可轉移的視覺模型》(https://arxiv.org/pdf/2103.00020)，提出了CLIP(對比語言圖像預訓練)，這是一個強大的深度學習模型，旨在以統一的方式理解和解釋圖像和文本。它結合了視覺和語言編碼器，將文本描述與視覺內容聯系起來。CLIP模型本身不生成圖像的描述，但可以用來評估文本和圖像之間的關系。例如，你可以提供一張貓的圖片，以及一個標簽列表，如“貓”和“狗”，以確定哪個標簽與圖片匹配的可能性最高。今天，這篇文章將涵蓋使用PyTorch從頭開始實現CLIP的過程。

CLIP(對比學習-圖像預訓練)

傳統的機器學習模型通常需要大量特定任務的標記數據集進行微調。例如，一個訓練用來識別狗的模型可能在識別貓方面表現不佳，除非它專門針對貓的圖片進行了微調。

CLIP的架構支持零樣本學習，這意味著它可以執行它沒有直接訓練過的任務，通過利用其在圖像和文本之間學到的廣泛關聯。例如，基于它們的文本描述，它可以對它在訓練期間從未見過的圖片進行分類。引用他們的論文：“我們在零樣本的情況下匹配原始ResNet-50在ImageNet上的準確性，而不需要使用它訓練時的128萬個訓練樣本。”

CLIP有以下我們需要構建的組件：

文本編碼器
圖像編碼器
自定義數據集(如果你正在訓練)
對稱損失

文本編碼器

由于我們的主要目標是使文本和視覺表示的嵌入對齊，我們將需要一個文本編碼器模型來為圖像的文本描述創建特征。本文不會涵蓋如何從頭開始構建文本編碼器，而是直接使用變換器庫來創建編碼器，盡管這將涵蓋CLIP實現的主要思想。為了簡單起見，使用Distil Bert模型是一個不錯的選擇，因為它輕量級，性能幾乎和標準BERT模型一樣好，具有類似的基礎架構。這是需要記住的一點，我們不是加載預訓練版本。

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'])

現在，TextEncoder的前向傳遞輸出將是(Batch_Size, Token_Size + 1, Embed_Size)，在標準BERT架構中，模型的目標是兩個任務，輸出一個額外的CLS_Token，附加到原始令牌前面，通常用于進一步微調分類任務，以及預測掩蔽令牌，使用掩蔽令牌前后的所有令牌的信息。

由于我們關心的是為我們的文本數據獲取特征嵌入，我們將只取[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)

層歸一化是深度學習中非常常見的概念，這不是我第一次解釋它，但讓我們再次解釋一下，我們有一個網絡的輸入，其中包含來自不同類別或特征的數據，因為在每個訓練周期中批次會變化，數據的分布也會變化，在一批中分布可能在[0, 2)范圍內，而在下一批中它可能有樣本分布在[0, 100]范圍內。在訓練過程中數據分布的變化被稱為協變量偏移。由于輸入的劇烈變化，輸出也會變化，損失也會變化，如果損失劇烈變化，那么在反向傳播過程中權重將以更高的幅度更新，導致梯度不平滑。簡而言之，歸一化輸入將限制其在整個訓練批次中的分布，因此，損失不會有劇烈變化，將導致更平滑的梯度和更快的訓練，幫助模型更多地關注學習特征。

圖像編碼器

CLIP有兩個圖像編碼器選項，ResNet或視覺變換器。我們已經開發了各種視覺變換器，因此將使用標準實現。如果你想使用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)

上面的編碼器類將圖像張量傳遞給模型，然后將其投影到與文本編碼器輸出相同的共同嵌入空間，后面是一個歸一化層。

自定義數據集

現在CLIP是一個(相當)密集的模型，所以如果你想從頭開始訓練它，你必須在一個小數據集上訓練它。由于本文只涉及如何從頭開始實現架構，我們將不會進一步詳細說明如何創建數據集，但為了示例，這可能是你想要做的。

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：你選擇的數據集中圖像的路徑列表。
texts：數據集中每張圖片的標題或文本描述。

自定義數據集類在調用Dataset類時創建分詞器并分詞所有文本，我們使用的是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()類是我們把所有放在一起的地方，架構將包括來自圖像和文本編碼器的嵌入，然后用于計算對稱損失。這是核心實現的NumPy樣式偽代碼。

在我們的實現中，我們將在CLIPModel類的前向函數中計算損失。第一步是獲取圖像和文本嵌入，然后進行交叉乘法以獲得相似性矩陣或logits。回到我們的第二張圖。

logits是通過取圖像和文本嵌入的點積來創建的，由于這篇論文基于對比學習，我們的主要目標是將文本表示與視覺對齊。那么計算相似性矩陣有什么幫助呢?

答案是每個從圖像編碼器接收到的圖像令牌(圖5：I_1,I_2,.., I_n; 其中I是嵌入，n是批次大小)乘以文本編碼器接收到的每個令牌。得到最終矩陣(B, Token, Embed)@(B, Embed, Token)→(B, Token, Token)。現在我們的任務是最大化每個對角線元素(I1T1, I2T2,…, InTn)的值。由于我們想要對齊我們的文本和視覺表示，相應的圖像令牌應該與其相應的文本最高相關。這就是我們將如何為批次中的所有圖像完成的，但讓我們看看單個令牌。

這里，圖并不是真的不同，我們取圖像嵌入I，并與批次中的每個文本嵌入計算點積。例如，當我們使用I3時，我們希望它與批次中相應的文本嵌入T3最強地對齊。理想情況下，I3行中最高的值應該是點積I3?T3，以同樣的方式批量處理，看起來就像我們在最大化所有對角線元素，其中每個In與其相應的Tn最佳對齊。為了實現這一點，我們使用一個損失函數來衡量每一行中最大值與其他值的突出程度。這實際上是通過取行和列的交叉熵損失來實現的。

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)
我們通過取點積來計算logits，如前所述，然后乘以溫度參數的指數。如果你熟悉對比學習或讀過我關于DINO(無標簽蒸餾)的文章，你可能知道通常使用除以溫度來銳化輸出分布。然而，我們不直接除以溫度，而是乘以一個可訓練的張量，該張量使用nn.Parameter()設置，并初始化為log(1/7)。由于eln(x)=x，那么exp(log(1/T))應該是1/T，你可能會想知道我們為什么不簡單地乘以1/T。原因是使用log(1/T)可以讓優化器在訓練期間更容易計算和更新梯度。這種方法在深度學習中是一種常見的做法，因為它可以帶來更平滑的訓練和更穩定的模型權重更新。
標簽簡單地用批次大小生成([0, 1,..N])。正如我們之前討論的，目標是最大化每個對角線元素(i1T1, i2T2,..inTn)，因此整個矩陣中每一行的標簽是[0, 1, 2, ..N]，對應于哪一行的元素應該是最大的。
如偽代碼所述，嵌入已經歸一化，但我們不需要這樣做，因為我們在返回圖像和文本編碼器的輸出時已經應用了層歸一化。
按照偽代碼，計算行和列的交叉熵損失。我們通過傳遞logits的轉置和正常的logits以及標簽，取兩個損失的平均值，現在我們有最終的損失結果。

設置模型

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

責任編輯：趙寧寧來源：小白玩轉Python