MIL-NCE与S3DG分布式训练:如何吃下HowTo100M海量视频-文本对 简介面向大规模视频理解的 MIL-NCE 方法完整实现基于 PyTorch 的 GPU 分布式训练代码解决千万级视频数据下高效训练与表示学习问题适合算法工程师、研究生及深度学习进阶者参考。压缩包共 22 个文件以 Python 脚本为主13 个 py涵盖模型结构定义、主训练脚本、多个数据集加载器与评估脚本并配有 5 个 csv 数据索引、说明文档、LICENSE 及日志文件整体仅 22.02MB目录结构清晰便于直接阅读和改造。已有 184 人学习下载。包内不仅提供 HowTo100M 视频列表还包含 YouCook2、MSRVTT、HMDB51 等下游任务验证集以及分布式训练入口可用于理解 torch.nn.parallel.DistributedDataParallel 的数据并行原理与多进程数据加载策略。借助这些代码可逐步完成从数据准备、单机多卡训练到下游任务评估的完整流程配套的说明文档、依赖清单与配置文件还能帮助快速安装依赖、设定学习率和批次大小并借助日志与检查点监控训练过程是实践大规模视频表示学习的优质参考。1. 百M级视频-文本对上的MIL-NCE与S3DG分布式训练落点HowTo100M上的视频-文本对存在大量弱对齐一段连续十几分钟的教学视频往往只有几十秒画面和字幕真正对应。这种数据规模下用标准softmax做对比学习分母面对百万级负样本根本算不动MIL-NCE的思路是把视频切成多个片段实例用噪声对比估计把归一化计算降成实例级二分类配合S3DG膨胀卷积从帧序列抽取时空特征。这套代码把HowTo100M预训练、PyTorch DDP多卡同步、HMDB51/YouCook2/MSRVTT下游评估串成整条链路main_distributed.py、loss.py、s3dg.py分别对应训练入口、损失函数和模型结构。适合正在做视频理解或跨模态检索、需要一份带数据并行和断点恢复的工程模板的工程师。拿到代码先别急着跑先想清楚三件事负样本怎么组织、数据怎么切分、checkpoint怎么恢复。2. MIL-NCE损失与S3DG结构为什么这一组合能吃下HowTo100M2.1 从NCE到MIL-NCE把softmax降成实例级判别HowTo100M有一百万个视频字幕由ASR自动生成和画面的对齐是稀疏且带噪的。早期做法是拿视频整体特征和文本特征做相似度再扔进softmax分母需要对所有候选视频做指数求和。视频规模到了千万级这个分母不可计算。NCE于是被引入不估计完整分布而是训练一个判别器去区分真实正样本和从噪声分布里采样出来的负样本。形式上把每个训练样本变成一个二分类问题负样本数量可以自由控制在几到几百计算量完全不依赖类别总数。MIL解决的是一整个视频只和少数文本片段相关的问题。如果把视频池化成单向量再和句子算相似度时间上不对齐的噪声会被平均进特征。更合理的做法是把视频分成若干片段每个片段得到独立的特征表示文本和所有片段分别算相似度再用max或attention聚合。训练信号因此能精确落在有语义回应的片段上不会被无关画面稀释。下面这段是模型里最核心的计算片段把视频的T个片段特征和batch内所有文本特征做内积在时间维上取max再对比# loss.py 里 MIL-NCE 的矩阵计算骨架 def mil_nce_loss(video_feats, text_feats, temperature0.07): # video_feats: [B, T, D] B个视频每个视频T个片段 # text_feats: [B, D] B条文本 sim torch.einsum(btd,nd-btn, video_feats, text_feats) sim sim / temperature # 时间维max选出最匹配的片段 sim_max, _ sim.max(dim1) # [B, B] labels torch.arange(B, devicesim.device) loss_v2t F.cross_entropy(sim_max, labels) loss_t2v F.cross_entropy(sim_max.t(), labels) return 0.5 * (loss_v2t loss_t2v)einsum把视频的片段维度和文本维度一次性做内积得到[B, T, B]的相似度张量再在时间维上取max变成[B, B]矩阵矩阵第i行第j列代表第i个视频和第j条文本的匹配分对角线是正样本其他位置是batch内负样本。temperature取0.05到0.2之间把分数缩放到对softmax梯度友好的区间取值过小时logits过于尖锐反向传播梯度直接消失。2.2 S3DG膨胀卷积和文本分支的取舍S3DG是把2D Inception的卷积核做时间维膨胀原本(k, k)的卷积展开成(k, k, k)输入是一段连续帧输出同时编码空间外观和短时运动。s3dg.py里通常的做法是用ImageNet预训练的2D权重初始化3D模型把权重沿时间维复制这是视频模型快速收敛的关键。输入侧video_loader.py按均匀间隔抽16或32帧空间分辨率224训练时随机裁剪加水平翻转验证时只做中心裁剪。文本分支没有直接上BERT。百万级别的视频对每条文本都过transformer计算量不现实常见做法是对ASR字幕做词级embedding平均再接一层线性投影映射到和视频特征相同的维度。词表由HowTo100M的字典裁剪得到低频词直接丢弃。这样的文本编码器在DDP训练里只占用很小的计算开销却已经能支撑起视频检索任务上有区分度的结果预训练结束后它还可以在下游HMDB51的zero-shot评测里充当类别原型编码器。2.3 为什么选DDP而不是DataParallelDataParallel是单进程多线程每次前向都要把输入从0号卡广播到其他卡梯度回传时再Reduce回0号卡通信频繁且Python线程的GIL会拖住数据装载。DDP每个rank是独立进程数据由DistributedSampler按rank切分成不相交的部分每个进程独立前向和反向只在optimizer.step之前对梯度做一次all-reduce通信量与batch大小解耦。这也解释了主训练脚本为什么叫main_distributed.py而不叫train.py——分布式训练入口和单卡版本本质上就是两套逻辑。3. main_distributed.py拆解DDP启动、数据切分与多卡训练循环拿到压缩包先看文件布局比逐行读代码更高效。main_distributed.py是训练入口args.py集中管理所有超参数s3dg.py定义模型loss.py定义损失video_loader.py、hmdb_loader.py、youcook_loader.py、msrvtt_loader.py分别负责不同数据集的加载eval_hmdb.py、eval_youcook.py、eval_msrvtt.py是三个下游评估脚本metrics.py提供Rk与MedR计算csv目录下是各数据集的划分清单checkpoint目录用来放中间权重log目录放训练日志。3.1 进程组初始化与torchrun启动参数DDP代码第一步是初始化进程组。main_distributed.py用的是env://模式环境变量由torchrun注入脚本内部不需要手动传递rank和world_sizeimport torch.distributed as dist def main_worker(rank, world_size, args): dist.init_process_group(backendnccl, init_methodenv://) torch.cuda.set_device(rank) model S3DG(args).cuda(rank) model DistributedDataParallel(model, device_ids[rank]) # 每个rank加载自己的数据分片之后进入训练循环init_process_group返回前会阻塞等待所有rank到齐只要一张卡启动失败整个训练卡在初始化这一步。遇到这种情况先查master_port是否被占用其次确认同一台机器上各进程的CUDA_VISIBLE_DEVICES设置一致。这和pytorch环境搭建时常见的多卡识别问题类似排查思路是一致的。torchrun的常用参数如下这些参数和init_process_group的env://模式是配套的参数作用常见值--nproc_per_node每台机器启动的进程数等于GPU个数--nnodes参与训练的机器数1或2--node_rank当前机器在集群中的编号从0开始--master_addrrank0所在机器的IP内网IP--master_port进程组通信端口29500# 单机4卡最常见 torchrun --nproc_per_node4 main_distributed.py --batch_size32 # 双机16卡每台机器分别执行 node1: torchrun --nnodes2 --nproc_per_node8 --node_rank0 \ --master_addr192.168.1.10 --master_port29500 main_distributed.py node2: torchrun --nnodes2 --nproc_per_node8 --node_rank1 \ --master_addr192.168.1.10 --master_port29500 main_distributed.py不必再用python -m torch.distributed.launchtorchrun把环境变量注入、子进程异常退出时的清理都做了少踩很多坑。master_addr只需要指向node0master_port选一个没有被防火墙拦掉的端口。进程组是否就绪可以分别打印rank和world_size两边能对上说明通信正常。3.2 DistributedSampler与视频解码的多进程配合args.py里几个参数要在启动前确认batch_size是每张卡的batch不是全局batchnum_workers控制每个DataLoader的加载线程数pin_memory建议打开。数据加载器用DistributedSampler切分样本在rank之间不重叠每个epoch必须手动调用set_epoch改变打乱顺序sampler DistributedSampler(dataset, num_replicasworld_size, rankrank) loader DataLoader(dataset, batch_sizeargs.batch_size, samplersampler, num_workers8, pin_memoryTrue) for epoch in range(args.epochs): sampler.set_epoch(epoch) for video, text in loader: video video.cuda(rank, non_blockingTrue) text text.cuda(rank, non_blockingTrue) # forward backward optimizer.step()set_epoch丢掉的后果是每个epoch样本顺序完全一致Dropout和BatchNorm这类依赖随机性的层学不到足够的扰动模型收敛变慢且更容易过拟合。另外用了DistributedSampler之后DataLoader的shuffle参数必须保持False否则会和采样器重复打乱数据分布反而被破坏。视频解码是分布式训练里最容易变成瓶颈的一环。cv2逐帧读视频在worker进程里非常慢常见做法是用decord的VideoReader批量解码后再采样worker数控制在4到8。同一块GPU上等数据的时间不该超过计算时间的20%一旦超过就去检查是不是每步都在重复开关视频文件描述符把视频帧先抽成jpg序列再喂DataLoader是治理这个瓶颈最直接的手段。3.3 checkpoint与会话恢复分布式训练的checkpoint统一由rank0写因为各rank模型参数在all-reduce后完全一致多写只会浪费磁盘。保存时取model.module.state_dict()而不是model.state_dict()否则多包一层DistributedDataParallel之后load阶段会出现key不匹配。恢复时map_location要指到当前rank的cuda设备ckpt torch.load(args.resume_from, map_locationfcuda:{rank}) model.module.load_state_dict(ckpt[model_state]) optimizer.load_state_dict(ckpt[optimizer_state]) sampler.load_state_dict(ckpt[sampler_state])恢复后loss值突然跳变基本是sampler状态没存。DistributedSampler内部记录了每个epoch的随机位移不恢复它重启后的数据顺序和中断那一刻对不上相当于拿另一个数据流继续训练指标自然回退。存储频率建议每5个epoch存一个同时把最近一个软链到last.pt断点恢复直接读last.pt即可。4. loss.py数值实现与HMDB51/YouCook2/MSRVTT评估脚本4.1 从相似度矩阵看MIL-NCE的负样本视野loss.py里的MIL-NCE矩阵骨架在2.1节展示过这里补充两个容易写错的细节。第一个是负样本来源当前batch内第i个视频和第i条文本是正样本对其余B-1个位置都是负样本这就是batch内对比学习。它的特点是负样本质量随batch增大而提升因为batch里出现语义相近样本的概率更高模型被迫学出更细的判别边界。负样本越多loss对相似度的梯度越集中在正样本附近模型越好区分相似样本单卡显存只允许batch8时负样本只有7个模型很快退化表现为loss不降、检索结果基本没有区分度。DDP下有个坑值得单独拿出来讲每个rank只在自己本卡上计算损失负样本总数是单卡batch-1和总卡数无关。想增大负样本视野可以在计算前用all_gather把所有rank的文本特征收集到一起再扩大相似度矩阵# 收集所有rank的文本特征作为额外负样本 text_list [torch.zeros_like(text_feats) for _ in range(world_size)] dist.all_gather(text_list, text_feats) text_all torch.cat(text_list, dim0) # [B * world_size, D] sim torch.einsum(btd,nd-btn, video_feats, text_all) labels torch.arange(B, devicesim.device) rank * B # 关键偏移注意最后一行label的偏移all_gather之后第i个视频的正样本文本位于全局索引rankBi处不偏移的话cross_entropy会把每张卡的0号样本当成当前样本负样本全部错位。代价是相似度矩阵从[B, B]变成[B, Bworld_size]显存和通信量都增加但每个rank的loss带了全局负样本梯度更稳收敛路径更接近大batch单卡的效果。4.2 梯度累积与优化器状态想让batch更大又不想爆显存可以上梯度累积。accumulation_steps必须配合no_sync上下文使用with model.no_sync(): loss.backward() # 每accumulation_steps步真正同步梯度 if step % accumulation_steps 0: for p in model.parameters(): p.grad / accumulation_steps optimizer.step()model.no_sync()让当前batch跳过DDP的梯度all-reduce否则每步都同步累积就失去意义。no_sync不能套在最后一个batch上否则最后一步的梯度不广播各rank参数开始分叉。优化器选型上百万级预训练常见做法是用SGD加momentum配合cosine退火微调阶段再切到AdamWargs.py里如果没有特别说明默认的weight_decay可以从1e-4起步过大的weight_decay在对比学习里会显著压碎特征范数。4.3 三个评估脚本zero-shot动作识别与双向检索eval_hmdb.py、eval_youcook.py、eval_msrvtt.py对应三个下游任务数据集加载器分别是hmdb_loader.py、youcook_loader.py、msrvtt_loader.pycsv目录下放着各自的划分文件。三个任务的评测方式完全不同看指标前先分清楚口径脚本数据集评测方式核心指标eval_hmdb.pyHMDB51zero-shot动作识别视频特征与类别文本特征最近邻Top-1, Top-5eval_youcook.pyYouCook2视频到文本检索R1, R5, MedReval_msrvtt.pyMSRVTT双向视频-文本检索R1, R5, R10HMDB51有51类动作每个类别用类别名套一个固定模板句子比如一个人在做xxx过文本编码器得到类别原型特征待测视频特征与51个原型算余弦距离取最近的类别作为预测。这是标准的zero-shot协议验证的是视频特征是否真正对齐了语义而不只是记忆标注。YouCook2的validation_youcook.csv是成对数据每个视频伴随多条描述句子测评时用视频特征去检索全部句子特征Rk表示正确答案出现在前k个结果中的比例MedR衡量排名中位数越小越好。metrics.py里recall_at_k计算有两个常见细节分数矩阵的形状是(query, candidate)方向搞反会得到完全错误的结果检索评估必须在全局样本上做不能用DDP各算各的否则每个rank只看到1/world_size的数据指标偏高且不可复现。用all_gather把特征collect到rank0再统一算指标虽然慢一点结果干净。4.4 数据加载器与csv的对应关系howto100m_videos.csv是预训练主数据all_videos.csv把多来源视频合并成一个清单。加载器读csv后按视频路径和文本对组织样本msrvtt_loader.py按msrvtt_test.csv的划分只保留测试子集YouCook2加载器读取validation_youcook.csvHMDB51加载器则负责把视频片段和类别标签对应起来。所有加载器最后统一返回视频张量、文本张量和可选的标签喂给训练和评估两个流程。5. 从预训练到下游迁移特征缓存、固定种子与温度系数排错5.1 特征缓存而不是端到端微调视频骨干网络参数量大HMDB51只有几千个样本直接端到端微调几乎必然过拟合。更稳的做法是用训练好的S3DG把所有下游视频过一遍把最后一层embedding存成npy后续只做线性分类或最近邻检索。eval脚本基本都走特征缓存路线先提特征再算指标显存占用小评估逻辑还能反复调。with torch.no_grad(): video_feats model(video_batch) # [B, D] np.save(ffeat_cache/{video_id}.npy, video_feats.cpu().numpy())缓存特征时保留一个视频id到特征路径的映射文件后面算R1要用它把特征和文本对应起来不要靠文件名排序去猜对应关系一旦csv顺序变化结果就错位。5.2 分布式训练固定随机种子复现实验时每个rank的seed要加rank偏移def set_seed(seed, rank): random.seed(seed rank) np.random.seed(seed rank) torch.manual_seed(seed rank) torch.cuda.manual_seed_all(seed rank)全部rank用同一个seed不会让数据重复但会让每个进程的数据增强完全同步负样本之间的扰动变小损失函数震荡幅度比预期更大。加上rank偏移后再配合DistributedSampler.set_epoch打乱和数据增强都能错开多卡训练曲线更平滑。5.3 训练发散时的三个检查点loss出现nan或者相似度矩阵长时间不变化不用急着改学习率按顺序排查temperature是否落在0.05到0.2区间太小的值会让softmax输出接近one-hot梯度消失视频和文本特征是否做了L2归一化没归一化时模长差异会主导相似度排序训练不稳定使用fp16混合精度时einsum在fp16下容易溢出把sim临时转回float32再参与除法。验证时可以先用单卡、batch32、固定seed跑50步观察loss是否平滑下降单卡正常而多卡发散重点检查all_gather后的label偏移和no_sync的使用位置多卡梯度不同步则检查梯度累积时no_sync有没有错误套在最后一个batch上。这三点排查完仍发散再看学习率warmup——视频模型前一个epoch从0.1倍学习率爬升到峰值比一上来就全速更稳定。本文还有配套的精品资源点击获取