ARTICLE · INTELLIGENCE

战地情报 · 详情页

来自尧图项目组的一线实战观察与深度解析

跨模态图文检索实战:双塔共享特征空间与对比学习全流程解析

跨模态图文检索实战:双塔共享特征空间与对比学习全流程解析 简介2024年“泰迪杯”数据挖掘挑战赛B题跨模态图文互检模型源码采用共享特征空间对比学习方案适合数据挖掘、人工智能、计算机视觉方向的学习者与竞赛选手用于理解图文语义关联和跨模态检索的实现路径。资源包共9个文件、仅10KB以4个Python脚本与2个YAML配置为主体辅以说明文档、依赖清单与Git忽略配置脚本覆盖数据处理、模型设计、训练与评估流程配置则对应训练和推理参数结构紧凑且便于逐模块研读。目前已有68人加入学习读者可对照完整代码梳理共享空间构建、对比学习目标设计、模型验证等关键环节并在此基础之上复现赛题思路、调整参数或迁移到相近的图文匹配任务中。同时项目说明与依赖清单能帮助快速搭建运行环境节省配置与调试时间。1. 2024 年泰迪杯 B 题基于共享特征空间的跨模态图文互检到底在解决什么2024年泰迪杯数据挖掘挑战赛B题的核心是基于共享特征空间对比学习的跨模态图文互检。我拆完这份源码的第一感受是竞赛题的难点从来不在模型结构本身而在把对比学习这套训练范式在图文数据上完整跑通并调稳。它解决的是真实检索问题——给一张图要从文本库里捞出匹配的描述给一句话要从图库里返回对应的图像评分看的是 RecallK 这类排序指标。源码覆盖了双塔编码、共享特征空间投影、InfoNCE 损失、负样本采样到评估的全流程。适合正在备赛的学生、想快速上手跨模态检索的工程师以及把竞赛代码当多模态入门教材的算法新人。2. 赛题拆解与模型选型B 题在考什么跨模态框架怎么定拿到一份竞赛源码先别急着分析训练代码把命题意图拆清楚。B 题给的是一批成对的图文数据模型要学的是判断“这张图和这句话是否匹配”并基于这种匹配能力完成双向检索。用大白话讲提交的模型要能让文本和图像在一个空间里直接比较相似度然后在候选库里按相似度排序。一旦进入排序问题模型的核心目标就从“分类正确”变成了“把正样本排到负样本前面”这和对比学习的优化目标完全一致。2.1 赛题在考什么从图文检索指标反推技术路线赛事评分一般用 RecallK 和 MRR 这两类指标。RecallK 的意思是在返回的前 K 个结果里是否包含正确匹配的那个样本。K 通常取 1、5、10Recall1 最严格要求第一个结果就命中Recall10 相对宽松只要前十个结果里出现就算得分。MRRMean Reciprocal Rank则关注正确样本排在第几位排得越靠前得分越高。从指标反推模型需要具备两个能力一是相似度度量能力匹配的图文相似度要明显高于不匹配的二是排序稳定性在几百上千条候选里正确样本要稳定落在靠前位置。这两个能力正好都是对比学习的强项。数据天然以“图像—文本”正样本对的形式组织不需要额外标注分类两个塔各自提取特征后直接用损失函数把正样本对拉近、负样本对推远就行。对比学习还有一层好处在竞赛时间有限的条件下尤其重要图像和文本的预训练权重都能直接拿开源模型做初始化训练时只需要微调。相比自己在单卡上从零训练一个大模型这种做法的起步成本低很多。所以整个 B 题的默认技术路线就是双塔结构 共享特征空间 对比损失。2.2 双塔结构与共享特征空间为什么检索场景默认不走单塔双塔结构的组成很直观图像塔用 ResNet 或 ViT文本塔用 BERT 或 RoBERTa两个塔各自把输入映射成特征向量。关键在后面的一层投影头Projection Head它通常是一到两层 MLP把图像特征和文本特征映射到同一个维度比如 256 维。只有映射到同一个维度之后图像和文本才能直接算余弦相似度。这里说的“共享特征空间”不是指两个编码器参数共享而是指两个塔输出的特征落在同一个向量坐标系里。投影头就是这个坐标系的控制器常见实现是 Linear 激活函数 Linear中间加 LayerNorm 效果更稳。投影头维度太小会丢失信息太大又会让损失函数难以约束128 到 512 是竞赛里比较常见的区间源码里默认的投影维度通常也落在这个范围。为什么不直接做一个单塔跨模态 Transformer单塔把图像 token 和文本 token 拼成一个序列送入 Transformer交互确实充分但检索场景有一个致命问题候选库里的每一条都要和 query 重新过一遍模型。库里有几万条文本就要跑几万次前向单塔拿到的特征无法离线缓存线上推理时每一对都要重新计算。双塔结构可以先把整库的特征离线算好存成向量文件线上推理时只需要算 query 这一侧的特征再做一次矩阵乘法就能得到全库排序。对比维度双塔结构单塔跨模态模型训练成本两个塔可并行显存占用可控长序列交互显存开销大推理速度库端特征预计算一次矩阵乘每对候选都要重新前向检索扩展性支持万级以上候选库候选库变大后几乎不可用特征交互深度浅层交互靠投影头对齐深层交互信息融合充分竞赛场景里数据量不大双塔的信息融合深度完全够用而评测时的检索效率是硬性的。源码里常见的做法是图像塔用 ResNet-50 级别的 backbone文本塔用中文 BERT两个塔都不大几万图文对的训练周期能控制在可接受范围内。2.3 对比学习目标函数InfoNCE 与温度系数InfoNCE 的思想可以用一句话概括在一个 batch 内部把正样本对区分出来。假设 batch 里有 N 对图文图像特征和文本特征相乘得到一个 N×N 的相似度矩阵第 i 行第 j 列表示第 i 张图和第 j 句文本的相似度。正样本就在对角线上每一行的目标是让对角线位置得分最高其余 N-1 个位置都是负样本。标准的对称对比损失会把这个多分类分别从两个方向各做一次第一个方向是“图像检索文本”每一行以文本为候选第二个方向是“文本检索图像”每一列以图像为候选。两个方向的交叉熵取平均就得到最终损失。这个实现方式在 CLIP 等模型中已经成为默认配置竞赛源码里通常也是这个结构少数变体只是换了温度系数的放置位置。温度系数 τ 是这里最敏感的超参数。相似度矩阵在进 softmax 之前会先除以 τ。τ 偏大时softmax 分布变得平滑负样本的梯度被稀释模型收敛稳定但判别力不足τ 偏小时softmax 分布尖锐模型被迫去区分最难区分的负样本但训练过程会变得激进容易出现震荡。常见默认值是 0.07对应 CLIP 的原始设置。竞赛实践中先按 0.07 启动loss 下降过慢再往 0.05 方向调出现震荡就往 0.1 方向退。共享特征空间不是天生对齐的而是被这个损失函数一步一步“逼”出来的。这也解释了为什么投影头维度、温度系数、负样本策略这些细节都会直接影响最终指标——它们都在改变对特征空间的约束方式。3. 数据预处理与负样本构建模型没跑起来之前一半的坑在数据侧双塔框架定下来之后真正费时间的其实是数据侧。跨模态任务的数据不像分类任务那样简单它需要同时保证图像和文本两个维度的质量还要设计正负样本的组织方式。我拆这份源码时最先做的就是数据扫描事实证明这一步帮后面省了大量排查时间。3.1 图文数据清洗先保证样本对干净竞赛数据包解压出来之后常见脏数据有三类。第一类是图像文件损坏下载不完整或编码异常训练时一读到这个文件就报错第二类是文本异常空值、重复、超长截断第三类是图文错标图是这张图、文本对应另一张图这类最隐蔽只能靠抽样人工检查。第一类和第二类可以脚本批量处理。from PIL import Image from pathlib import Path def scan_images(img_dir): bad_list [] for path in Path(img_dir).rglob(*.jpg): try: with Image.open(path) as im: im.verify() # 校验文件完整性不触发完整解码 except Exception: bad_list.append(str(path)) return bad_list这段代码的逻辑很简单但有两个细节值得注意。im.verify()只读取文件头并校验解码信息不会把整个图像加载进内存扫描一个大目录时速度快得多。捕获Exception而不是只捕获OSError是因为截断文件在 verify 阶段可能抛出多种异常类型统一捕获更稳妥。拿到损坏列表之后从标注文件里把这些样本整体剔除对应的文本记录也要一起删否则后面 Dataset 里会出现“图不存在”的索引错误。文本清洗相对琐碎但同样必要。英文描述统一转小写、折叠多余空格、去掉 HTML 标签中文描述没有大小写问题但要清理控制字符和全角半角混乱。之后按 image_id 做去重防止同一条文本被重复训练导致模型对个别样本过拟合。长度过滤也在这里做太短的描述信息量不足太长的描述可能包含噪声一般设置在 2 到 64 个字符区间。3.2 Dataset 实现图像与文本预处理对齐预处理的关键在于对齐。图像侧的 Resize、裁剪、归一化文本侧的 tokenize、padding、截断要在同一个 Dataset 里完成保证每个 step 返回的都是严格对应的张量对。图像用 TorchVision 的 transform文本用 HuggingFace 的 tokenizer两边各管各的最后在__getitem__里拼在一起返回。from torch.utils.data import Dataset from transformers import AutoTokenizer from torchvision import transforms class ITMDataset(Dataset): def __init__(self, df, img_dir, max_len64): self.df df self.img_dir img_dir self.tokenizer AutoTokenizer.from_pretrained(bert-base-chinese) self.tf transforms.Compose([ transforms.Resize((224, 224)), transforms.RandomResizedCrop(224, scale(0.8, 1.0)), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]), ]) self.max_len max_len def __getitem__(self, idx): row self.df.iloc[idx] img Image.open(f{self.img_dir}/{row[image_name]}).convert(RGB) text row[text] tokens self.tokenizer( text, paddingmax_length, truncationTrue, max_lengthself.max_len, return_tensorspt) return (self.tf(img), tokens[input_ids].squeeze(0), tokens[attention_mask].squeeze(0))__getitem__返回三个张量处理后图像、input_ids 和 attention_mask。文本侧在 tokenizer 内部已经做了 padding 到 max_length所以 DataLoader 的默认 collate_fn 就能直接工作不需要额外写 collate 逻辑这一点能省掉不少初学者经常会犯的维度错误。参数推荐值说明图像尺寸224×224竞赛数据量不大224 够用随机裁剪 scale0.81.0保留主体避免裁掉关键目标文本最大长度64常见图文描述都在此范围内归一化ImageNet 均值方差使用预训练权重的硬性要求数据增强要克制。分类任务里常用的 ColorJitter 在图文互检里风险很高描述里有“金黄色的麦田”时颜色扰动直接破坏语义对应关系。我一般只保留 RandomResizedCrop 和随机水平翻转让模型不依赖目标绝对位置但不改动颜色信息。3.3 负样本采样策略in-batch 与 hard negative 并用对比学习的负样本质量决定训练上限。in-batch negatives 是最简单的方案batch 里 N 对图文对每行来说其他 N-1 个样本都是负样本不需要额外采样计算效率高。但它的局限也很明显batch 里大部分负样本和正样本差异巨大模型很容易学会“图像颜色完全不同不匹配”这种低层次判别之后指标提升就停滞了。常见改进做法是在 in-batch 之外追加 hard negative。所谓 hard negative就是“很像但不对应”的负样本。对图文互检来说两个文本描述同一类物体但对应不同图像这两条文本互为 hard negative两张图构图相近但物体不同也互为 hard negative。def sample_hard_negatives(embedding_matrix, positive_indices, k2): sim embedding_matrix embedding_matrix.T hard_pairs [] for i in range(len(positive_indices)): sim[i, positive_indices[i]] -float(inf) # 屏蔽正样本 hard_idx sim[i].topk(k).indices.tolist() hard_pairs.append((i, hard_idx)) return hard_pairs这段代码的输入是某个模态的特征矩阵输出是每个样本最相似但不对应的 topk 索引。屏蔽正样本那一步是关键如果不把对角线位置置为负无穷topk 第一个选出来的往往是它自己。实际工程里不会每个 step 都重算全库特征常见做法是每训练几个 epoch用当前模型重新编码一次训练子集更新 hard negative 列表。hard negative 的比例要控制。比例太低起不到提难作用比例太高模型会被噪声支配。常见做法是把 hard negative 控制在 batch 内负样本数量的 10% 到 20%更新频率每隔 3 到 5 个 epoch 一次。这个参数和温度系数一样属于最后阶段的精调对象基础模型还没跑稳之前不建议动。4. 训练流程与参数配置把对比损失跑稳跑低的实操笔记训练部分的代码是整个源码包的核心也是调参空间最大的地方。按“损失实现→训练循环→评估脚本”的顺序拆开讲每一段代码都是可以直接抄进工程的最小可运行版本。4.1 InfoNCE 损失实现对比损失的核心实现不到十行难点在于理解它为什么这样写。图像塔输出图像特征文本塔输出文本特征两个特征都先做 L2 归一化然后矩阵相乘得到相似度矩阵。除以温度系数后对每一行做交叉熵。对称版本会再把矩阵转置算一次文本到图像方向两个损失取平均。import torch import torch.nn.functional as F def contrastive_loss(img_feats, txt_feats, temp0.07): img_feats F.normalize(img_feats, dim-1) txt_feats F.normalize(txt_feats, dim-1) logits img_feats txt_feats.T / temp labels torch.arange(len(img_feats), deviceimg_feats.device) loss_i2t F.cross_entropy(logits, labels) loss_t2i F.cross_entropy(logits.T, labels) return (loss_i2t loss_t2i) / 2F.normalize放在矩阵相乘之前这一步不是可选项。如果不做归一化特征向量的模长差异会干扰温度系数的语义logits 的绝对值尺度不稳定调 τ 就变成了猜谜。labels 用torch.arange生成正好对应矩阵对角线位置每一行的正样本都是它自己。温度系数的位置也值得注意temp出现在矩阵相乘之后、softmax 之前。有个常见误解有人会把temp乘到 normalize 之前效果完全不同。除温度再 softmax 是 CLIP 的标准写法它会等比放大所有 logits锐化概率分布让 loss 更关注相似度最高的那些负样本。4.2 训练循环与优化器参数训练循环和常规分类训练差别不大真正的门道在梯度累积和梯度裁剪上。对比损失的梯度幅度对 batch 组成敏感偶尔一个 hard negative 样本就能让梯度激增所以裁剪几乎是必须的。accum_steps 2 # 等效大batch的累积步数 total_steps (len(train_loader) * epochs) // accum_steps # OneCycleLR 总步数 optimizer torch.optim.AdamW(model.parameters(), lr5e-5, weight_decay0.02) scheduler torch.optim.lr_scheduler.OneCycleLR( optimizer, max_lr5e-5, total_stepstotal_steps, pct_start0.05) for epoch in range(epochs): for step, (img, input_ids, attn_mask) in enumerate(train_loader): img, input_ids, attn_mask img.cuda(), input_ids.cuda(), attn_mask.cuda() img_feats, txt_feats model(img, input_ids, attn_mask) loss contrastive_loss(img_feats, txt_feats, temp0.07) (loss / accum_steps).backward() if (step 1) % accum_steps 0: torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() scheduler.step() optimizer.zero_grad()梯度累积的实现方式是把 loss 除以 accum_steps 再 backward这样累积 accum_steps 步之后梯度均值等于一个大 batch 的梯度学习率不用为小 batch 额外调整。clip_grad_norm_的 max_norm 设为 1.0在对比学习里是很稳的选择尤其在 hard negative 比例较高时能挡住梯度爆炸。参数推荐值调整方向学习率5e-5backbone 冻结底层时可放宽到 1e-4weight_decay0.02投影头过拟合时加大到 0.05温度系数0.07loss 不降往 0.05震荡往 0.1梯度累积24根据显存反向决定梯度裁剪1.0稳定优先不改注意预训练权重是竞赛的底线。两个塔都从开源预训练权重起步loss 在前几个 epoch 就能掉到 1 以下从零训练的话几万图文对的数据量根本喂不饱 ResNet 和 BERT效果会差一个量级。4.3 评估脚本RecallK 与 MRR评估不能在线逐条调用模型正确做法是把验证集完整的特征先算一遍全部收集起来然后一次性做矩阵运算。验证集通常不会太大一个 epoch 内存足够容纳。def evaluate(img_feats, txt_feats, k_list(1, 5, 10)): img_feats F.normalize(img_feats, dim-1) txt_feats F.normalize(txt_feats, dim-1) sim txt_feats img_feats.T # 文本检索图像方向 for k in k_list: topk sim.topk(k, dim-1).indices correct (topk torch.arange(len(sim))[:, None]).any(dim-1) recall correct.float().mean().item() print(ft2i Recall{k}: {recall:.4f})sim矩阵的行是文本列是图像sim[i][j] 表示第 i 句文本和第 j 张图的相似度。topk取前 k 个图像的索引当某一行的 topk 索引里包含行号本身时说明正确样本被召回。MRR 的实现同理找到第一个正确索引的位置后取倒数。提示评估脚本里的归一化和训练脚本必须逐行一致否则相似度排序会失真。loss 曲线很漂亮但指标完全对不上多半是这里出的问题。有了评估脚本之后整个训练闭环就完整了。跑通这一遍接下来才谈得上优化指标。5. 避坑记录跨模态训练中五个花了不少时间才定位的翻车现场复现这份源码的完整流程走下来模型结构本身几乎没有卡住人反而是各种“看似合理但实际上在跨模态场景下会出问题”的默认配置一个接一个地消耗时间。下面五个坑都来自实际操作每条按“现象 → 原因 → 解决”写清楚。5.1 三个“看起来正常但指标不对”的典型第一个坑训练 loss 正常下降但验证集 Recall1 始终是 0现象loss 从 4.x 一路降到 1.x曲线非常好看跑到验证阶段却完全不行top5 偶尔能蒙对top1 永远命中不了。原因训练时contrastive_loss内部对特征做了 L2 归一化而评估脚本直接用了模型的原始输出。两个特征空间的尺度不一致余弦相似度的排序自然失真正样本在原始特征空间里并没有被拉到应有的位置。解决评估脚本里在计算相似度矩阵之前对两侧特征做同样的F.normalize。从那以后我养成了一个检查习惯每次写评估脚本都会和训练函数逐行对比预处理逻辑而不是想当然地认为模型输出可以直接用。第二个坑双向指标差距悬殊text2img 远好于 img2text现象同一个模型、同一份特征text2img 的 Recall5 能到 40%img2text 只有 20%差距接近一倍。原因图像塔和文本塔的收敛速度不一致。BERT 文本塔的输出空间本身语义性强从预训练权重起步后很快就能对齐ResNet 图像塔的输出偏底层视觉特征需要更多 epoch 才能被对比损失拉到语义空间。解决一是给图像塔的投影头输出维度从 256 提到 512让图像特征有更多表达空间二是加强图像侧的 RandomResizedCrop 比例迫使图像塔关注语义主体而非背景细节。本质上是让弱的塔承担更多优化压力。第三个坑加了 hard negative 之后 loss 反而震荡现象hard negative 比例从 0 调整到 20% 之后训练 loss 出现周期性震荡验证指标不升反降。原因hard negative 列表是用上一个 epoch 的特征采样出来的而特征空间本身在训练中持续移动等到当前 epoch 真正用到这些负样本时监督信号已经部分失效。模型面对的是一批“旧空间里的难样本”强行拟合只会引入噪声。解决降低更新频率每 5 个 epoch 重采样一次 hard negative并把比例控制在 10% 左右。如果仍然震荡说明是难负样本的噪声在主导梯度继续下调比例或者对 hard negative 样本的 loss 降权处理。5.2 两个“一改就崩”的参数陷阱第四个坑显存不够把 batch size 从 128 降到 32指标全面下滑现象loss 变大验证指标掉了 5 个点以上而且不只是 Recall1所有 K 值都在掉。原因对比学习对 batch 内负样本数量有硬依赖。batch 越小负样本越少模型每个 step 能看到的信息量大幅缩水梯度估计的方差也会变大。这不是学习率没调的问题是方法本身的约束。解决不要直接缩小 batch用梯度累积模拟大 batch。实测用 batch32 加 accum4等效于 128 的 batch指标能拉回九成。但这里还有一个隐藏问题如果 backbone 带 BatchNorm梯度累积不会改变 BatchNorm 统计量的计算方式它仍然只基于 32 个样本更新统计量需要换成 SyncBN 才能在等效大 batch 下正常工作。第五个坑复现结果不稳定同样的参数跑两次差几个点现象同样的代码、同样的参数连续跑两次验证集指标波动超过 3 个百分点有时候波动大到影响超参数判断。原因直接原因是随机种子没有固定数据加载顺序和图像增强的随机性没有被约束更深层的原因是对比学习对负样本分布极其敏感shuffle 顺序不同导致每个 batch 的负样本构成不同模型在每个 step 接收的监督信号就有差异。解决固定 torch、numpy、random 三个库的随机种子并且给 DataLoader 传worker_init_fn保证数据加载线程的随机性也可控。固定种子之后实验方差能压到 1 个点以内扫描超参数才有可比性。6. 进阶用法用消融实验把“玄学”变成可复现的参数baseline 跑通之后拉开差距的是消融实验。对比学习在竞赛里常被说成玄学但把几个关键参数分开扫一遍它其实每个参数的影响都可解释。我通常把温度系数、投影头维度、hard negative 比例作为三个扫描变量其他配置保持不变。{ temp: 0.07, proj_dim: 256, hard_neg_ratio: 0.1, lr: 5e-5, batch_size: 64 }for t in 0.03 0.07 0.1; do python train.py --temp $t --proj_dim 256 --hard_neg_ratio 0.1 done每一组实验的验证结果都要单独记录不要只看最优值。温度从 0.07 降到 0.05 时如果 Recall1 上升但 Recall5 下滑说明模型把注意力过度集中在难负样本上排序前几位变准但整体排序变毛糙了。投影头维度从 256 加到 512 时如果 text2img 方向的提升明显大于 img2text说明图像塔还在学习阶段投影头多出来的维度给了它追赶的空间。验证阶段除了算指标拆这份源码时我习惯把验证集特征保存成 npy 文件写一个最近邻检索脚本人工抽查 top10 检索结果。典型 query 分三类描述物体位置的、描述颜色的、描述动作的。如果“穿红色外套的男人”检索结果里出现“黑色外套的女人”颜色语义没有对齐如果出现构图相似但物体完全不同的图说明视觉底层特征在主导排序文本语义没有真正参与。从那以后我每次拿到竞赛源码都强制自己先把最小配置跑干净确认 loss 和验证指标能对齐再动任何超参数。宁可多花半天把数据清洗和评估脚本做严格也不要在训练跑完一轮之后才发现评估有 bug。对比学习的每个参数都有迹可循前提是你给它一个可复现的评测环境。希望帮到你。本文还有配套的精品资源点击获取
RELATED READING

延伸阅读

更多一线实战笔记与深度复盘,助您持续精进