跨模态图文互检实战：共享特征空间对比学习源码解析与调优-开发者社区

简介：这份资源是2024年“泰迪杯”数据挖掘挑战赛B题的完整参赛源码，面向数据挖掘、人工智能与计算机视觉方向的高校学生及竞赛选手，聚焦跨模态图文互检这一典型任务。方案以共享特征空间对比学习为核心思路，通过对图文特征进行对齐与对比约束，建模两种模态之间的语义关联，从而支撑高效的图文互检。压缩包共9个文件，约10KB，以4个Python脚本为主体，覆盖数据处理、模型定义、训练与评估等环节，另含2个yaml配置文件用于管理训练与评估参数，以及说明文档和依赖清单，结构紧凑、便于快速复现。目前已有68人学习。读者可借此理解赛题从数据组织、模型搭建到训练优化的完整实现路径，掌握对比学习在跨模态检索中的应用方式，并参考配置与评估脚本进行实验调参与结果验证，适合作为课程设计、竞赛复现与相关研究的入门参考，仅供学习使用。

1. 从赛题到落地：这套跨模态图文互检源码到底能跑出什么结果

如果你正在准备数据挖掘类竞赛，或者手头有一个“以图搜文、以文搜图”的需求，这套 2024 年“泰迪杯”B 题的完整源码值得你花一个下午拆一遍。它解决的核心问题很具体：给一张图，从候选文本库里找出最匹配的描述；给一段文本，从候选图片库里找出最匹配的图。不是简单的标签分类，而是跨模态检索——两个模态的特征要在同一个空间里可比。源码里用的是共享特征空间对比学习，这个思路在近两年的多模态检索赛题里反复出现，属于主流且可复现的基线方案。适合谁？打过一两次数据挖掘比赛、能读懂 PyTorch 训练循环、想找一个结构完整的跨模态项目练手的人。如果你连 DataLoader 都没写过，建议先补基础，否则调试成本会很高。

2. 共享特征空间对比学习：为什么这样设计，以及数据怎么喂进去

2.1 跨模态检索的核心矛盾与对比学习的切入方式

跨模态检索最朴素的做法是：图片过 CNN 提特征，文本过 LSTM 或 BERT 提特征，然后各自算相似度。但这样做有个致命问题——图像特征和文本特征各自分布在不同的向量空间里，余弦相似度算出来的值没有可比性。共享特征空间要做的就是：把两个模态的特征映射到同一个维度、同一个分布区间里，让“匹配的图文对”在这个空间里距离近，“不匹配的图文对”距离远。

对比学习恰好是干这个的。它的训练逻辑是：一个 batch 里有 N 对正样本（图 + 对应文本），通过数据增强或负采样构造出负样本对，然后让正样本对的相似度尽可能高，负样本对的相似度尽可能低。源码里用的是 InfoNCE 损失的一个变体，温度系数 τ 控制分布的尖锐程度。这个设计的好处是：不需要显式的类别标签，只需要“哪段文本对应哪张图”这种弱监督信号，非常适合竞赛数据集的标注形式。

我拆源码时注意到一个细节：图像侧用的是 ResNet-50 去掉最后的全连接层，输出 2048 维特征后接一个投影头降到 256 维；文本侧用的是 BiLSTM + 注意力池化，同样投影到 256 维。两个投影头结构对称，但参数不共享。这是常见做法——共享的是最终的特征空间，不是中间层的权重。如果你想让模型更轻，可以把 ResNet-50 换成 ResNet-18，但检索精度会掉 3 到 5 个点，竞赛场景下不建议省这个算力。

2.2 数据加载与图文对构造：Dataset 和 Collate_fn 怎么写

竞赛给的数据通常是“图片文件夹 + 文本描述文件”的形式，每张图对应一段或多段描述。源码里的 Dataset 类做了三件事：读图、读文本、把图文对拼成一个样本。关键点在于文本长度不固定，不能直接 stack，需要自定义 collate_fn 做 padding。

import torch from torch.utils.data import Dataset, DataLoader from PIL import Image import pandas as pd class CrossModalDataset(Dataset): def __init__(self, img_dir, caption_file, transform=None, tokenizer=None, max_len=32): self.img_dir = img_dir self.df = pd.read_csv(caption_file) # 列：image_id, caption self.transform = transform self.tokenizer = tokenizer self.max_len = max_len def __len__(self): return len(self.df) def __getitem__(self, idx): row = self.df.iloc[idx] img_path = f"{self.img_dir}/{row['image_id']}.jpg" image = Image.open(img_path).convert("RGB") if self.transform: image = self.transform(image) # 文本转 id 序列，截断或填充到 max_len tokens = self.tokenizer(row['caption'], max_len=self.max_len) return image, torch.tensor(tokens, dtype=torch.long) def collate_fn(batch): images, texts = zip(*batch) images = torch.stack(images, dim=0) texts = torch.stack(texts, dim=0) # 已在 Dataset 里统一长度 return images, texts

这段代码的逻辑说明：__getitem__返回的是单对图文，collate_fn负责把一个 batch 里的样本拼成张量。参数max_len=32是文本截断长度，竞赛里描述通常不超过 20 个词，32 够用。如果你的数据集描述更长，改到 64，但显存占用会线性增加。tokenizer需要自己实现一个简单的词表映射，源码里用的是按词频取 top-10000 构建词表，OOV 统一映射到<UNK>。

注意：collate_fn里不要用pad_sequence再转 tensor，那样会多一次拷贝。直接在 Dataset 里 padding 好，collate 只做 stack，训练速度会快 10% 左右。

2.3 模型前向与损失计算：投影头和温度系数怎么调

模型部分分三个模块：图像编码器、文本编码器、共享投影头。图像编码器用 torchvision 的 resnet50，文本编码器用 nn.LSTM 加一个注意力层。投影头是两层 MLP，中间加 ReLU 和 Dropout。

import torch.nn as nn import torchvision.models as models class ImageEncoder(nn.Module): def __init__(self, out_dim=256): super().__init__() backbone = models.resnet50(pretrained=True) self.backbone = nn.Sequential(*list(backbone.children())[:-1]) # 去掉 fc self.proj = nn.Sequential( nn.Linear(2048, 512), nn.ReLU(), nn.Dropout(0.3), nn.Linear(512, out_dim) ) def forward(self, x): feat = self.backbone(x).flatten(1) # (B, 2048) return self.proj(feat) class TextEncoder(nn.Module): def __init__(self, vocab_size, embed_dim=300, hidden_dim=512, out_dim=256): super().__init__() self.embed = nn.Embedding(vocab_size, embed_dim, padding_idx=0) self.lstm = nn.LSTM(embed_dim, hidden_dim, batch_first=True, bidirectional=True) self.attn = nn.Linear(hidden_dim * 2, 1) self.proj = nn.Sequential( nn.Linear(hidden_dim * 2, 512), nn.ReLU(), nn.Dropout(0.3), nn.Linear(512, out_dim) ) def forward(self, x): emb = self.embed(x) # (B, L, 300) out, _ = self.lstm(emb) # (B, L, 1024) weights = torch.softmax(self.attn(out), dim=1) # (B, L, 1) pooled = (out * weights).sum(dim=1) # (B, 1024) return self.proj(pooled)

参数说明：out_dim=256是共享空间的维度，源码里试过 128 和 512，256 在验证集上 Recall@1 最高。Dropout(0.3)是防止过拟合，竞赛数据量通常不大，这个值不要低于 0.2。文本侧的hidden_dim=512配合双向 LSTM，输出 1024 维，再经过注意力池化压成 1024 维向量。注意力池化的作用是让模型自己决定哪些词更重要，比直接取最后一个时间步的输出更稳。

损失函数用 InfoNCE，核心是构造相似度矩阵，对角线是正样本对，其余是负样本。温度系数 τ 初始设 0.07，训练过程中可以固定也可以学习。源码里是固定值，我建议你先固定跑通，再尝试可学习温度。

def info_nce_loss(img_feat, txt_feat, temperature=0.07): # 归一化到单位球面 img_feat = nn.functional.normalize(img_feat, dim=1) txt_feat = nn.functional.normalize(txt_feat, dim=1) logits = img_feat @ txt_feat.t() / temperature # (B, B) labels = torch.arange(logits.size(0), device=logits.device) loss_i2t = nn.functional.cross_entropy(logits, labels) loss_t2i = nn.functional.cross_entropy(logits.t(), labels) return (loss_i2t + loss_t2i) / 2

逻辑说明：logits是 B×B 的相似度矩阵，第 i 行第 j 列表示第 i 张图和第 j 段文本的相似度。对角线是正样本，所以labels就是[0, 1, 2, ..., B-1]。loss_i2t是图到文的检索损失，loss_t2i是文到图的检索损失，两者取平均。温度系数越小，模型对难负样本的惩罚越大，但太小会导致训练不稳定。0.07 是 SimCLR 里的经典值，跨模态场景下也适用。

3. 训练流程与评估指标：从跑通到跑出有竞争力的结果

3.1 训练循环与学习率调度

训练循环本身不复杂，但有几个竞赛场景下的细节要注意。源码里用的是 AdamW 优化器，初始学习率 1e-4，权重衰减 1e-4。图像编码器的 backbone 学习率是投影头的 0.1 倍，这是为了避免预训练权重被过快破坏。学习率调度用余弦退火，每 10 个 epoch 重启一次。

import torch.optim as optim from torch.optim.lr_scheduler import CosineAnnealingWarmRestarts # 分组学习率 backbone_params = list(model.img_encoder.backbone.parameters()) other_params = [p for n, p in model.named_parameters() if 'backbone' not in n] optimizer = optim.AdamW([ {'params': backbone_params, 'lr': 1e-5}, {'params': other_params, 'lr': 1e-4} ], weight_decay=1e-4) scheduler = CosineAnnealingWarmRestarts(optimizer, T_0=10, T_mult=2) for epoch in range(50): model.train() for images, texts in train_loader: images, texts = images.cuda(), texts.cuda() img_feat = model.img_encoder(images) txt_feat = model.txt_encoder(texts) loss = info_nce_loss(img_feat, txt_feat) optimizer.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) optimizer.step() scheduler.step()

参数说明：T_0=10表示第一次重启周期是 10 个 epoch，T_mult=2表示每次重启后周期翻倍。clip_grad_norm_的max_norm=1.0是防止梯度爆炸，对比学习里梯度突然变大很常见，这个操作能救你一命。batch size 源码里设的是 64，如果你的显存不够，降到 32，但温度系数要相应调大一点，比如 0.1，因为负样本数量少了。

3.2 评估指标：Recall@K 和 mAP 怎么算

跨模态检索的评估指标主要是 Recall@K 和 mAP。Recall@K 的意思是：对于每个查询，检索结果的前 K 个里有没有正确答案。mAP 是平均精度均值，考虑了正确答案的排名位置。源码里两个都实现了，竞赛提交通常看 Recall@1 和 Recall@5。

def evaluate(model, dataloader, k_list=[1, 5, 10]): model.eval() all_img_feats, all_txt_feats = [], [] with torch.no_grad(): for images, texts in dataloader: images, texts = images.cuda(), texts.cuda() all_img_feats.append(model.img_encoder(images)) all_txt_feats.append(model.txt_encoder(texts)) img_feats = torch.cat(all_img_feats, dim=0) txt_feats = torch.cat(all_txt_feats, dim=0) img_feats = nn.functional.normalize(img_feats, dim=1) txt_feats = nn.functional.normalize(txt_feats, dim=1) sim_matrix = img_feats @ txt_feats.t() # (N, N) ranks = torch.argsort(sim_matrix, dim=1, descending=True) labels = torch.arange(sim_matrix.size(0)).unsqueeze(1) recall = {} for k in k_list: correct = (ranks[:, :k] == labels).any(dim=1).float().mean().item() recall[f'R@{k}'] = correct return recall

逻辑说明：先把所有样本的特征算出来，拼成矩阵，然后算相似度矩阵。argsort得到每个查询的检索排名，labels是正确答案的索引。(ranks[:, :k] == labels).any(dim=1)判断前 k 个里有没有正确答案。这个评估方式假设每个查询只有一个正确答案，如果竞赛允许多个正确答案，需要改成is_in判断。源码里还实现了 mAP，逻辑类似，但计算的是每个正确结果的精度平均值。

注意：评估时一定要把模型切到eval()模式，并且用torch.no_grad()，否则显存会爆。另外，如果测试集很大，相似度矩阵可能放不下，需要分块计算。

3.3 推理与结果提交：如何生成竞赛要求的格式

竞赛提交通常要求一个 CSV 文件，每行是“查询 ID, 检索结果 ID, 排名”。源码里提供了一个生成脚本，核心逻辑是：对每个查询，取相似度最高的前 K 个结果，按排名写入。

import pandas as pd def generate_submission(model, query_loader, gallery_loader, k=10): model.eval() query_feats, query_ids = [], [] gallery_feats, gallery_ids = [], [] with torch.no_grad(): for images, texts, ids in query_loader: query_feats.append(model.img_encoder(images.cuda())) query_ids.extend(ids) for images, texts, ids in gallery_loader: gallery_feats.append(model.txt_encoder(texts.cuda())) gallery_ids.extend(ids) query_feats = nn.functional.normalize(torch.cat(query_feats), dim=1) gallery_feats = nn.functional.normalize(torch.cat(gallery_feats), dim=1) sim_matrix = query_feats @ gallery_feats.t() topk = torch.topk(sim_matrix, k, dim=1) rows = [] for i, qid in enumerate(query_ids): for rank, idx in enumerate(topk.indices[i]): rows.append({'query_id': qid, 'result_id': gallery_ids[idx], 'rank': rank + 1}) pd.DataFrame(rows).to_csv('submission.csv', index=False)

参数说明：k=10是提交的检索结果数量，竞赛通常要求 10 或 20。query_loader和gallery_loader的 batch size 可以设大一点，因为不需要反向传播。如果查询集和图库集是同一个数据集，可以合并计算，但要注意去重。源码里还处理了“查询结果不能是自身”的逻辑，如果你的竞赛允许自身匹配，可以去掉这个限制。

4. 避坑与排查：这套源码跑起来最容易翻车的五个地方

4.1 现象：训练 loss 不下降，一直卡在 4.5 左右

原因：最常见的是文本 tokenizer 的 padding_idx 没设对。如果 padding 的 id 不是 0，Embedding 层会把 padding 当成正常词来学，导致文本特征被噪声淹没。另一个可能是图像预处理用了 ImageNet 的均值和方差，但你的数据集分布差异很大。

解决：检查nn.Embedding的padding_idx是否和 tokenizer 里<PAD>的 id 一致。源码里<PAD>的 id 是 0，所以padding_idx=0。如果数据集是灰度图或医学图像，把归一化的均值和方差改成数据集自己的统计值，用torchvision.transforms.Normalize(mean=[0.5], std=[0.5])先跑通。

4.2 现象：验证集 Recall@1 比训练集低 20 个点

原因：过拟合。竞赛数据量通常只有几千对，ResNet-50 参数量太大，很容易记住训练集。源码里虽然加了 Dropout，但不够。

解决：把图像编码器的 backbone 冻结前 30 个 epoch，只训练投影头。或者用更强的数据增强：随机裁剪、颜色抖动、随机擦除。文本侧可以加词替换增强，把同义词随机替换。如果还不行，把 ResNet-50 换成 ResNet-18，参数量少一半，泛化会好很多。

4.3 现象：训练到一半 loss 突然变成 NaN

原因：梯度爆炸。对比学习的 loss 对温度系数很敏感，τ 太小或者学习率太大都会导致梯度爆炸。另外，如果 batch 里有全 padding 的文本，LSTM 的输出会全零，归一化时除以零也会产生 NaN。

解决：加梯度裁剪，clip_grad_norm_(model.parameters(), max_norm=1.0)。温度系数不要低于 0.05。检查数据里有没有空文本，如果有，在 Dataset 里过滤掉或者用<UNK>填充。源码里在 collate_fn 里加了一个断言，如果文本全 padding 就重新采样。

4.4 现象：评估时显存爆了，batch size 降到 1 还是爆

原因：评估时把所有样本的特征拼成一个大矩阵算相似度，如果测试集有 10000 个样本，相似度矩阵就是 10000×10000，float32 下要占 400MB，加上中间变量很容易爆。

解决：分块计算相似度。每次取 1000 个查询，和全部图库算相似度，取 topk 后释放中间变量。源码里提供了一个chunked_evaluate函数，逻辑就是两层循环，外层遍历查询块，内层遍历图库块。

def chunked_similarity(query_feats, gallery_feats, chunk_size=1000): results = [] for i in range(0, query_feats.size(0), chunk_size): q_chunk = query_feats[i:i+chunk_size] sim_chunk = q_chunk @ gallery_feats.t() results.append(sim_chunk) return torch.cat(results, dim=0)

4.5 现象：提交的 CSV 文件排名全是乱的，和本地评估结果对不上

原因：查询 ID 和图库 ID 的映射关系搞错了。竞赛给的 ID 可能是字符串，不是连续的整数，如果直接用torch.arange生成 labels，评估时会对不上。

解决：在 Dataset 里把原始 ID 存下来，评估和生成提交时都用原始 ID 做映射。源码里用了一个id2idx和idx2id的字典，确保顺序一致。另外，提交前用 pandas 读一遍 CSV，检查行数和格式是否符合竞赛要求。

5. 进阶技巧：把 Recall@1 再往上推 3 个点的几个实操手段

5.1 难负样本挖掘：让模型在“像但不是”的样本上多学

对比学习的瓶颈往往在于负样本太容易区分。随机采样的负样本里，大部分和正样本差异很大，模型学不到细粒度的区分能力。难负样本挖掘的思路是：在每个 epoch 结束后，用当前模型算一遍相似度，找出那些“相似度高但不是正确答案”的样本对，下一个 epoch 重点训练这些对。

具体做法：维护一个难负样本队列，每次取 batch 时，从队列里采样一部分作为额外的负样本。源码里实现了一个简化版：每个 epoch 结束后，对每个查询取 top-10 检索结果，如果正确答案不在前 10，就把这些错误结果加入难负样本池。下一个 epoch 的 loss 计算时，把这些难负样本的相似度乘以一个惩罚系数。

def hard_negative_loss(img_feat, txt_feat, hard_neg_indices, temperature=0.07, penalty=2.0): img_feat = nn.functional.normalize(img_feat, dim=1) txt_feat = nn.functional.normalize(txt_feat, dim=1) logits = img_feat @ txt_feat.t() / temperature # 对难负样本加大惩罚 for i, neg_idx in enumerate(hard_neg_indices): if neg_idx >= 0: logits[i, neg_idx] *= penalty labels = torch.arange(logits.size(0), device=logits.device) loss_i2t = nn.functional.cross_entropy(logits, labels) loss_t2i = nn.functional.cross_entropy(logits.t(), labels) return (loss_i2t + loss_t2i) / 2

参数说明：penalty=2.0是惩罚系数，不要设太大，否则模型会过度关注难负样本，导致对简单样本的区分能力下降。hard_neg_indices是一个长度为 batch size 的数组，每个元素是难负样本在 batch 里的索引，如果没有就填 -1。这个技巧在竞赛后期能带来 1 到 2 个点的提升，但训练时间会增加 30% 左右。

5.2 模型集成：两个不同随机种子的模型怎么融合

单模型的 Recall@1 到 60 左右就很容易过拟合，集成是性价比最高的提升手段。源码里提供了两种集成方式：特征拼接和相似度平均。特征拼接是把两个模型的特征 concat 起来，再算相似度；相似度平均是分别算相似度矩阵，然后取平均。

我一般用相似度平均，因为不需要重新训练，直接加载两个模型权重，分别推理，最后平均。注意两个模型的投影维度要一致，否则没法平均。如果维度不同，可以先做 PCA 降到相同维度。集成两个模型通常能涨 2 到 3 个点，集成三个以上收益递减，而且推理时间线性增加。

def ensemble_similarity(model1, model2, query_loader, gallery_loader): sim1 = compute_similarity(model1, query_loader, gallery_loader) sim2 = compute_similarity(model2, query_loader, gallery_loader) return (sim1 + sim2) / 2

5.3 后处理：相似度归一化和 top-k 重排

推理阶段有一个容易被忽略的技巧：对相似度矩阵做行归一化。因为不同查询的相似度分布可能不同，有的查询所有相似度都偏高，有的都偏低，直接排序会不公平。行归一化后，每个查询的相似度变成相对值，排序更合理。

另一个技巧是 top-k 重排：先取前 50 个候选，然后用一个更精细的模型（比如交叉注意力）对这 50 个重新打分。源码里没有实现重排，但你可以加一个简单的 MLP 重排器，输入是查询特征和候选特征的拼接，输出是一个分数。这个重排器可以用少量数据训练，推理时只对 top-50 算，时间开销可接受。