跨模态检索实战：基于CLIP模型的对比学习架构优化与性能调优

1次阅读

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

跨模态检索在电商搜索、内容审核等场景中有着广泛的需求。比如在电商平台，用户可能用文字描述搜索商品图片，或者反过来用图片搜索相关商品描述。传统方法通常面临以下问题：

特征空间不一致：图像和文本特征通常由不同模型提取，导致特征空间不匹配
人工标注成本高：需要大量人工标注的图文对数据
计算效率低下：传统的跨模态检索方法通常需要复杂的特征融合和匹配过程

CLIP 模型采用双编码器设计，分别处理图像和文本输入，并通过对比学习实现模态对齐。架构如下：

[Image Input]       [Text Input]
    |                   |
[ViT Encoder]      [Text Encoder]
    |                   |
[Projection]       [Projection]
    \                 /
     \               /
[对比损失 (InfoNCE)]

关键组件说明：

ViT- 图像编码器：基于 Vision Transformer 架构，将图像分割为 patch 序列进行处理
文本编码器：通常使用 Transformer 架构处理文本序列
Projection Head：将不同模态的特征映射到同一语义空间
InfoNCE 损失：通过负样本对比学习实现模态对齐

from torchvision import transforms
from transformers import AutoTokenizer

# 图像预处理
train_transform = transforms.Compose([transforms.RandomResizedCrop(224),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.48145466, 0.4578275, 0.40821073],
        std=[0.26862954, 0.26130258, 0.27577711]
    )
])

# 文本预处理
tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
texts = ["a photo of a cat", "a picture of a dog"]
inputs = tokenizer(texts, padding=True, return_tensors="pt")

import torch.nn as nn

class ProjectionHead(nn.Module):
    def __init__(self, input_dim, output_dim=256):
        super().__init__()
        self.projection = nn.Linear(input_dim, output_dim)
        self.gelu = nn.GELU()
        self.fc = nn.Linear(output_dim, output_dim)

    def forward(self, x):
        x = self.projection(x)
        x = self.gelu(x)
        x = self.fc(x)
        return x

def contrastive_loss(logits, temperature=0.07):
    # logits: (batch_size, batch_size)
    labels = torch.arange(logits.size(0), device=logits.device)
    loss_i = F.cross_entropy(logits/temperature, labels)
    loss_t = F.cross_entropy(logits.t()/temperature, labels)
    return (loss_i + loss_t) / 2

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

with autocast():
    image_features = image_encoder(images)
    text_features = text_encoder(input_ids, attention_mask)
    loss = contrastive_loss(image_features, text_features)

scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

accumulation_steps = 4

for i, (images, texts) in enumerate(train_loader):
    # Forward pass
    loss = model(images, texts)

    # Normalize loss
    loss = loss / accumulation_steps

    # Backward pass
    loss.backward()

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

torch.onnx.export(
    model,
    (dummy_image, dummy_text),
    "clip_model.onnx",
    input_names=["image", "text"],
    output_names=["features"],
    dynamic_axes={"image": {0: "batch_size"},
        "text": {0: "batch_size"},
        "features": {0: "batch_size"}
    },
    opset_version=13
)