基于clap对比学习预训练模型的文本-图像跨模态检索实战

1次阅读

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

跨模态检索的核心挑战在于文本和图像特征空间的对齐。传统方法存在明显局限性：

关键词匹配方法（如 BM25）：仅依赖表面词汇匹配，无法理解深层语义
CNN+RNN 组合模型：
特征提取与对齐分步进行，存在误差累积
模态间交互仅通过后期融合实现，难以捕捉细粒度关联

这些方法在 COCO 数据集上的 Recall@10 通常不超过 40%，表明传统方案已遇到性能瓶颈。

CLAP 模型通过 InfoNCE 损失函数实现特征对齐：

# 简化版 InfoNCE 实现
class ContrastiveLoss(nn.Module):
    def __init__(self, temp=0.07):
        super().__init__()
        self.temp = temp  # 温度系数控制样本区分度

    def forward(self, z_img, z_text):
        # 计算相似度矩阵
        logits = (z_text @ z_img.T) / self.temp
        labels = torch.arange(len(z_img)).to(device)
        loss = F.cross_entropy(logits, labels) + F.cross_entropy(logits.T, labels)
        return loss/2

温度系数的选择直接影响模型性能：
– 值过大：所有样本相似度趋同
– 值过小：导致梯度爆炸

class CLAP(nn.Module):
    def __init__(self):
        super().__init__()
        self.text_encoder = RobertaModel.from_pretrained('roberta-base')
        self.img_encoder = ViTModel.from_pretrained('google/vit-base-patch16-224')
        # 投影头将不同模态特征映射到统一空间
        self.text_proj = nn.Linear(768, 256)
        self.img_proj = nn.Linear(768, 256)

关键设计要点：
– 文本塔使用 RoBERTa 处理变长文本
– 图像塔采用 ViT 提取 patch 特征
– 投影头维度需保持一致（通常 256-512 维）

In-batch 负采样：
利用同批次样本天然构成负对
计算高效但可能包含假阴性样本
难样本挖掘：
定期筛选高相似度的错误匹配对
需配合指数移动平均 (EMA) 模型判断

class CrossModalDataset(Dataset):
    def __init__(self, df, img_dir, tokenizer, img_size=224):
        self.df = df
        self.tokenizer = tokenizer
        self.transform = transforms.Compose([transforms.Resize((img_size, img_size)),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ])

    def __getitem__(self, idx):
        row = self.df.iloc[idx]
        img = Image.open(row['image_path']).convert('RGB')
        return {'image': self.transform(img),
            'input_ids': self.tokenizer(row['text'], 
                                      padding='max_length', 
                                      max_length=77,
                                      truncation=True)['input_ids']
        }

混合精度训练配置

scaler = torch.cuda.amp.GradScaler()

with autocast():
    img_feat = model.img_encoder(images)
    text_feat = model.text_encoder(input_ids)
    loss = criterion(img_feat, text_feat)

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

学习率调度策略

scheduler = get_cosine_schedule_with_warmup(
    optimizer,
    num_warmup_steps=1000,  # 逐渐增大学习率
    num_training_steps=total_steps
)

使用 DistributedDataParallel 替代DataParallel
同步 BN 保证批次统计量一致性

model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
model = DDP(model, device_ids=[local_rank])

使用 Faiss 建立向量索引

index = faiss.IndexIVFPQ(faiss.IndexFlatIP(256),  # 内积相似度
    dimension=256,
    nlist=100,
    M=8,  # 压缩维度
    nbits_per_idx=8
)
index.train(image_features)
index.add(image_features)

量化部署

quantized_model = torch.quantization.quantize_dynamic(model, {nn.Linear}, dtype=torch.qint8
)

文本截断长度需与预训练模型一致（RoBERTa 通常 77token）
图像裁剪比例建议不低于原图的 80%

# 特征可视化
features = torch.cat([img_feats, text_feats])
tsne = TSNE(n_components=2)
vis_feats = tsne.fit_transform(features.cpu())

# 绘制不同模态特征分布
plt.scatter(vis_feats[:batch_size,0], vis_feats[:batch_size,1], label='Image')
plt.scatter(vis_feats[batch_size:,0], vis_feats[batch_size:,1], label='Text')

在 Flickr30k 测试集上的表现：