使用 PyTorch 从头开始构建 CLIP | 对比语言图像预训练-51CTO.COM

在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