简介这份资源是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, transformNone, tokenizerNone, max_len32): 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_lenself.max_len) return image, torch.tensor(tokens, dtypetorch.long) def collate_fn(batch): images, texts zip(*batch) images torch.stack(images, dim0) texts torch.stack(texts, dim0) # 已在 Dataset 里统一长度 return images, texts这段代码的逻辑说明__getitem__返回的是单对图文collate_fn负责把一个 batch 里的样本拼成张量。参数max_len32是文本截断长度竞赛里描述通常不超过 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_dim256): super().__init__() backbone models.resnet50(pretrainedTrue) 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_dim300, hidden_dim512, out_dim256): super().__init__() self.embed nn.Embedding(vocab_size, embed_dim, padding_idx0) self.lstm nn.LSTM(embed_dim, hidden_dim, batch_firstTrue, bidirectionalTrue) 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), dim1) # (B, L, 1) pooled (out * weights).sum(dim1) # (B, 1024) return self.proj(pooled)参数说明out_dim256是共享空间的维度源码里试过 128 和 512256 在验证集上 Recall1 最高。Dropout(0.3)是防止过拟合竞赛数据量通常不大这个值不要低于 0.2。文本侧的hidden_dim512配合双向 LSTM输出 1024 维再经过注意力池化压成 1024 维向量。注意力池化的作用是让模型自己决定哪些词更重要比直接取最后一个时间步的输出更稳。损失函数用 InfoNCE核心是构造相似度矩阵对角线是正样本对其余是负样本。温度系数 τ 初始设 0.07训练过程中可以固定也可以学习。源码里是固定值我建议你先固定跑通再尝试可学习温度。def info_nce_loss(img_feat, txt_feat, temperature0.07): # 归一化到单位球面 img_feat nn.functional.normalize(img_feat, dim1) txt_feat nn.functional.normalize(txt_feat, dim1) logits img_feat txt_feat.t() / temperature # (B, B) labels torch.arange(logits.size(0), devicelogits.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_decay1e-4) scheduler CosineAnnealingWarmRestarts(optimizer, T_010, T_mult2) 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_norm1.0) optimizer.step() scheduler.step()参数说明T_010表示第一次重启周期是 10 个 epochT_mult2表示每次重启后周期翻倍。clip_grad_norm_的max_norm1.0是防止梯度爆炸对比学习里梯度突然变大很常见这个操作能救你一命。batch size 源码里设的是 64如果你的显存不够降到 32但温度系数要相应调大一点比如 0.1因为负样本数量少了。3.2 评估指标RecallK 和 mAP 怎么算跨模态检索的评估指标主要是 RecallK 和 mAP。RecallK 的意思是对于每个查询检索结果的前 K 个里有没有正确答案。mAP 是平均精度均值考虑了正确答案的排名位置。源码里两个都实现了竞赛提交通常看 Recall1 和 Recall5。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, dim0) txt_feats torch.cat(all_txt_feats, dim0) img_feats nn.functional.normalize(img_feats, dim1) txt_feats nn.functional.normalize(txt_feats, dim1) sim_matrix img_feats txt_feats.t() # (N, N) ranks torch.argsort(sim_matrix, dim1, descendingTrue) labels torch.arange(sim_matrix.size(0)).unsqueeze(1) recall {} for k in k_list: correct (ranks[:, :k] labels).any(dim1).float().mean().item() recall[fR{k}] correct return recall逻辑说明先把所有样本的特征算出来拼成矩阵然后算相似度矩阵。argsort得到每个查询的检索排名labels是正确答案的索引。(ranks[:, :k] labels).any(dim1)判断前 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, k10): 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), dim1) gallery_feats nn.functional.normalize(torch.cat(gallery_feats), dim1) sim_matrix query_feats gallery_feats.t() topk torch.topk(sim_matrix, k, dim1) 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, indexFalse)参数说明k10是提交的检索结果数量竞赛通常要求 10 或 20。query_loader和gallery_loader的 batch size 可以设大一点因为不需要反向传播。如果查询集和图库集是同一个数据集可以合并计算但要注意去重。源码里还处理了“查询结果不能是自身”的逻辑如果你的竞赛允许自身匹配可以去掉这个限制。4. 避坑与排查这套源码跑起来最容易翻车的五个地方4.1 现象训练 loss 不下降一直卡在 4.5 左右原因最常见的是文本 tokenizer 的 padding_idx 没设对。如果 padding 的 id 不是 0Embedding 层会把 padding 当成正常词来学导致文本特征被噪声淹没。另一个可能是图像预处理用了 ImageNet 的均值和方差但你的数据集分布差异很大。解决检查nn.Embedding的padding_idx是否和 tokenizer 里PAD的 id 一致。源码里PAD的 id 是 0所以padding_idx0。如果数据集是灰度图或医学图像把归一化的均值和方差改成数据集自己的统计值用torchvision.transforms.Normalize(mean[0.5], std[0.5])先跑通。4.2 现象验证集 Recall1 比训练集低 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_norm1.0)。温度系数不要低于 0.05。检查数据里有没有空文本如果有在 Dataset 里过滤掉或者用UNK填充。源码里在 collate_fn 里加了一个断言如果文本全 padding 就重新采样。4.4 现象评估时显存爆了batch size 降到 1 还是爆原因评估时把所有样本的特征拼成一个大矩阵算相似度如果测试集有 10000 个样本相似度矩阵就是 10000×10000float32 下要占 400MB加上中间变量很容易爆。解决分块计算相似度。每次取 1000 个查询和全部图库算相似度取 topk 后释放中间变量。源码里提供了一个chunked_evaluate函数逻辑就是两层循环外层遍历查询块内层遍历图库块。def chunked_similarity(query_feats, gallery_feats, chunk_size1000): results [] for i in range(0, query_feats.size(0), chunk_size): q_chunk query_feats[i:ichunk_size] sim_chunk q_chunk gallery_feats.t() results.append(sim_chunk) return torch.cat(results, dim0)4.5 现象提交的 CSV 文件排名全是乱的和本地评估结果对不上原因查询 ID 和图库 ID 的映射关系搞错了。竞赛给的 ID 可能是字符串不是连续的整数如果直接用torch.arange生成 labels评估时会对不上。解决在 Dataset 里把原始 ID 存下来评估和生成提交时都用原始 ID 做映射。源码里用了一个id2idx和idx2id的字典确保顺序一致。另外提交前用 pandas 读一遍 CSV检查行数和格式是否符合竞赛要求。5. 进阶技巧把 Recall1 再往上推 3 个点的几个实操手段5.1 难负样本挖掘让模型在“像但不是”的样本上多学对比学习的瓶颈往往在于负样本太容易区分。随机采样的负样本里大部分和正样本差异很大模型学不到细粒度的区分能力。难负样本挖掘的思路是在每个 epoch 结束后用当前模型算一遍相似度找出那些“相似度高但不是正确答案”的样本对下一个 epoch 重点训练这些对。具体做法维护一个难负样本队列每次取 batch 时从队列里采样一部分作为额外的负样本。源码里实现了一个简化版每个 epoch 结束后对每个查询取 top-10 检索结果如果正确答案不在前 10就把这些错误结果加入难负样本池。下一个 epoch 的 loss 计算时把这些难负样本的相似度乘以一个惩罚系数。def hard_negative_loss(img_feat, txt_feat, hard_neg_indices, temperature0.07, penalty2.0): img_feat nn.functional.normalize(img_feat, dim1) txt_feat nn.functional.normalize(txt_feat, dim1) 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), devicelogits.device) loss_i2t nn.functional.cross_entropy(logits, labels) loss_t2i nn.functional.cross_entropy(logits.t(), labels) return (loss_i2t loss_t2i) / 2参数说明penalty2.0是惩罚系数不要设太大否则模型会过度关注难负样本导致对简单样本的区分能力下降。hard_neg_indices是一个长度为 batch size 的数组每个元素是难负样本在 batch 里的索引如果没有就填 -1。这个技巧在竞赛后期能带来 1 到 2 个点的提升但训练时间会增加 30% 左右。5.2 模型集成两个不同随机种子的模型怎么融合单模型的 Recall1 到 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) / 25.3 后处理相似度归一化和 top-k 重排推理阶段有一个容易被忽略的技巧对相似度矩阵做行归一化。因为不同查询的相似度分布可能不同有的查询所有相似度都偏高有的都偏低直接排序会不公平。行归一化后每个查询的相似度变成相对值排序更合理。另一个技巧是 top-k 重排先取前 50 个候选然后用一个更精细的模型比如交叉注意力对这 50 个重新打分。源码里没有实现重排但你可以加一个简单的 MLP 重排器输入是查询特征和候选特征的拼接输出是一个分数。这个重排器可以用少量数据训练推理时只对 top-50 算时间开销可接受。注意后处理不要过度调参竞赛的测试集分布和验证集可能不同过度拟合验证集的后处理参数在测试集上会翻车。我一般只做行归一化重排器只在验证集上确认有稳定提升才用。从那以后我每次跑跨模态检索的实验都会先把基线跑通确认 Recall1 在合理范围内再逐步加技巧。每加一个技巧都要在验证集上确认提升是稳定的而不是随机波动。这套源码的结构很清晰改起来不费劲但坑也不少希望帮到你。本文还有配套的精品资源点击获取