CLIP模型预训练实战：从零构建跨模态理解系统

1次阅读

共计 4179 个字符，预计需要花费 11 分钟才能阅读完成。

跨模态学习的目标是让模型能够理解不同模态（如图像和文本）之间的关联。CLIP（Contrastive Language-Image Pretraining）是 OpenAI 提出的一种跨模态预训练模型，它通过对比学习的方式，将图像和文本映射到同一个语义空间，从而实现高效的跨模态检索。

CLIP 的核心思想是使用双塔结构，分别对图像和文本进行编码，然后通过对比损失函数来优化这两个编码器。具体来说，CLIP 包含以下组件：

图像编码器：通常使用 ResNet 或 ViT（Vision Transformer）架构，将图像转换为固定维度的向量。
文本编码器：通常使用 Transformer 架构，将文本转换为相同维度的向量。
对比损失函数：通过计算图像和文本向量之间的相似度，优化模型使其能够正确匹配对应的图像 - 文本对。

数学上，对比损失函数可以表示为：

$$
\mathcal{L} = -\frac{1}{N} \sum_{i=1}^N \left[\log \frac{e^{\text{sim}(I_i, T_i)/\tau}}{\sum_{j=1}^N e^{\text{sim}(I_i, T_j)/\tau}} + \log \frac{e^{\text{sim}(T_i, I_i)/\tau}}{\sum_{j=1}^N e^{\text{sim}(T_i, I_j)/\tau}} \right]
$$

其中，$I_i$ 和 $T_i$ 分别表示第 $i$ 个图像和文本的向量表示，$\text{sim}$ 是相似度函数（通常为余弦相似度），$\tau$ 是温度参数，用于控制分布的尖锐程度。

在跨模态学习中，双编码器和交叉注意力是两种常见的架构。它们的优缺点如下：

双编码器（如 CLIP）：
优点：计算效率高，适合大规模预训练；推理速度快，适合实时应用。
缺点：模态间的交互较弱，可能丢失一些细粒度的关联信息。
交叉注意力架构（如 UNITER）：
优点：模态间交互强，能够捕捉更细粒度的关联。
缺点：计算复杂度高，训练和推理速度较慢；显存占用大，难以扩展到大规模数据。

对于大多数工业级应用，双编码器架构（如 CLIP）是更实用的选择，因为它在性能和效率之间取得了较好的平衡。

以下是 CLIP 模型的 PyTorch 实现代码，包含数据加载、对比损失计算和梯度累积等关键模块。

import torch
from torch.utils.data import Dataset, DataLoader
from PIL import Image

class CLIPDataset(Dataset):
    def __init__(self, image_paths, texts, transform=None):
        self.image_paths = image_paths
        self.texts = texts
        self.transform = transform

    def __len__(self):
        return len(self.image_paths)

    def __getitem__(self, idx):
        image = Image.open(self.image_paths[idx]).convert('RGB')
        if self.transform:
            image = self.transform(image)
        text = self.texts[idx]
        return image, text

import torch.nn as nn
from transformers import BertModel

class ImageEncoder(nn.Module):
    def __init__(self, model_name='resnet50'):
        super().__init__()
        self.model = torch.hub.load('pytorch/vision', model_name, pretrained=True)
        self.model.fc = nn.Identity()  # Remove the final classification layer

    def forward(self, x):
        return self.model(x)

class TextEncoder(nn.Module):
    def __init__(self, model_name='bert-base-uncased'):
        super().__init__()
        self.model = BertModel.from_pretrained(model_name)

    def forward(self, input_ids, attention_mask):
        outputs = self.model(input_ids=input_ids, attention_mask=attention_mask)
        return outputs.last_hidden_state[:, 0, :]  # Use the [CLS] token representation

class CLIPModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.image_encoder = ImageEncoder()
        self.text_encoder = TextEncoder()

    def forward(self, images, input_ids, attention_mask):
        image_features = self.image_encoder(images)
        text_features = self.text_encoder(input_ids, attention_mask)
        return image_features, text_features

def contrastive_loss(image_features, text_features, temperature=0.07):
    # Normalize features
    image_features = image_features / image_features.norm(dim=1, keepdim=True)
    text_features = text_features / text_features.norm(dim=1, keepdim=True)

    # Compute similarity matrix
    logits = (image_features @ text_features.T) / temperature

    # Labels are the indices of the diagonal elements
    labels = torch.arange(logits.shape[0], device=image_features.device)

    # Cross entropy loss
    loss_i = nn.functional.cross_entropy(logits, labels)
    loss_t = nn.functional.cross_entropy(logits.T, labels)
    loss = (loss_i + loss_t) / 2
    return loss

model = CLIPModel().cuda()
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
accumulation_steps = 4

for epoch in range(10):
    for i, (images, texts) in enumerate(train_loader):
        images = images.cuda()
        input_ids = tokenize(texts).cuda()
        attention_mask = (input_ids != 0).cuda()

        image_features, text_features = model(images, input_ids, attention_mask)
        loss = contrastive_loss(image_features, text_features)

        # Normalize the loss to account for gradient accumulation
        loss = loss / accumulation_steps
        loss.backward()

        if (i + 1) % accumulation_steps == 0:
            optimizer.step()
            optimizer.zero_grad()

在大规模预训练中，显存是宝贵的资源。以下是一些显存优化的技巧：

梯度检查点（Gradient Checkpointing）：通过牺牲部分计算时间换取显存节省。PyTorch 中可以通过 torch.utils.checkpoint 实现。
混合精度训练 ：使用torch.cuda.amp 进行自动混合精度训练，可以显著减少显存占用并加速训练。
模型并行：将模型的不同部分放置在不同的 GPU 上，适用于超大规模模型。
梯度累积：如上所示，通过累积多个小批量的梯度再进行一次参数更新，可以有效减少显存占用。