InfoXLM Fairseq 数据模块详解Dataset、Dictionary 与批次迭代器的设计与实战【免费下载链接】unilmLarge-scale Self-supervised Pre-training Across Tasks, Languages, and Modalities项目地址: https://gitcode.com/GitHub_Trending/un/unilm本篇基于 InfoXLM 仓库内置的 fairseq 数据模块文档 data.rst 展开系统讲解fairseq.data的四大核心——基础 Dataset 抽象、核心双语/单语数据集、辅助包装数据集Helper Datasets、词典Dictionary与批次迭代器Iterators。读完后你将掌握如何读懂 fairseq 数据集的collater/num_tokens/ordered_indices三件套如何复用ConcatDataset、RoundRobinZipDatasets等包装器组合多语语料以及EpochBatchIterator的分片sharding、可复现 shuffle 与断点恢复机制在源码层面是如何实现的。一、fairseq.data模块全景文档定义了哪些核心构件data.rst 以 Sphinx autoclass 方式组织将数据模块划分为四个层次这与仓库中 fairseq/data/__init__.py 的导出结构一一对应文档章节文档列出的类源码位置Datasets基础数据集FairseqDataset、LanguagePairDataset、MonolingualDatasetfairseq_dataset.py、language_pair_dataset.py、monolingual_dataset.pyHelper Datasets辅助数据集BacktranslationDataset、ConcatDataset、ResamplingDataset、RoundRobinZipDatasets、TransformEosDataset同目录下的 backtranslation_dataset.py 等对应文件DictionaryDictionarydictionary.pyIterators迭代器CountingIterator、EpochBatchIterator、GroupedIterator、ShardedIteratoriterators.py文档对数据集的定义是“Datasets define the data format and provide helpers for creating mini-batches”数据集定义数据格式并提供创建 mini-batch 的辅助工具。这正对应了FairseqDataset上collater方法的职责。此外__init__.py还导出了大量文档未逐一展开的实用包装器如PadDataset、RollDataset、MaskTokensDataset、SubsampleDataset等它们遵循同一套接口约定可作为深入阅读的延伸。二、FairseqDataset数据集契约与批处理辅助接口所有数据集都继承自 FairseqDataset它本身是torch.utils.data.Dataset与EpochListening混入类的组合定义了 fairseq 数据集必须遵守的“契约”__getitem__(index)/__len__()标准 PyTorch Dataset 接口必须实现collater(samples)将一批样本合并为一个可直接喂给模型的 mini-batchdict是整个数据模块的核心num_tokens(index)返回样本的 token 数用于批处理时强制执行--max-tokens约束size(index)返回样本大小用于--max-positions过滤ordered_indices()返回一个有序索引列表批次将按该顺序构建默认实现就是np.arange(len(self))即自然顺序supports_prefetch/prefetch(indices)可选的预取能力声明后训练器可在 epoch 开始前批量预取数据set_epoch(epoch)来自EpochListening混入类第 10-15 行每当 epoch 递增时接收通知默认空实现供ResamplingDataset等需要按 epoch 重新采样的数据集覆写。同文件还定义了 FairseqIterableDataset面向流式或无法在单机上随机索引的数据只要求实现__iter__。理解这组接口的关键在于分工“取单条样本”__getitem__、“定批次顺序”ordered_indices、“拼成张量”collater三件事解耦这使得各种包装数据集只需覆写其中一部分即可复用其余逻辑。三、核心数据集一LanguagePairDataset与平行语料批处理LanguagePairDataset 是机器翻译任务的平行语料数据集它本身只是两个 Dataset源、目标加尺寸数组的组合但collater承担了几乎全部批处理逻辑。3.1 构造参数与默认值从 构造函数第 149-180 行 可以确认文档 autoclass 背后的完整参数表参数默认值说明src/src_sizes/src_dict必填源数据集、源句长度列表、源词汇表tgt/tgt_sizes/tgt_dictNone目标侧可为空纯源语料场景left_pad_sourceTrue源句左填充符合自回归编码器从左到右的注意力习惯left_pad_targetFalse目标句右填充max_source_positions/max_target_positions1024/1024句长上限shuffleTrue批处理前是否打乱input_feedingTrue是否为 teacher forcing 生成移位后的prev_output_tokensremove_eos_from_source/append_eos_to_targetFalse/False用于源/目标互换反方向翻译时的 EOS 修正align_datasetNone词对齐数据集提供时必须同时有源和目标append_bosFalse是否在源/目标句首追加 BOS源码中还有一个值得注意的硬约束当提供tgt_dict时会断言源、目标词典的pad()、eos()、unk()索引必须一致第 159-162 行这保证了两个词表共享同一组特殊符号索引。3.2collate的输出结构与批内排序collate 函数第 12-113 行 生成的 batch 包含id样本在原始输入中的 IDnsentences/ntokens批内句数与目标侧总 token 数无目标侧时为源侧总 token 数net_input.src_tokens/src_lengths填充后的源张量与真实长度net_input.prev_output_tokens目标右移一位的张量input_feedingTrue时通过move_eos_to_beginning将 EOS 移到句首实现target用于计算损失的目标张量若提供词对齐则额外输出alignments与align_weights按目标索引出现频率取逆频加权的对齐权重。值得注意的是collate 内部会按源长降序重排 batch 内样本第 49-53 行配合左填充可以显著降低实际注意力计算量。3.3ordered_indices先随机后按长度稳定排序ordered_indices第 265-274 行 的实现是“随机打乱 → 按目标长度mergesort稳定排序 → 再按源长度稳定排序”。这意味着批与批之间尽量长度相近减少 padding 浪费而同一长度组内部保持随机序。num_tokens返回max(src_sizes, tgt_sizes)size返回(src_len, tgt_len)元组分别服务于--max-tokens与--max-positions。四、核心数据集二MonolingualDataset与语言模型目标构造MonolingualDataset 面向单语数据语言建模、BPE 学习等。它最有价值的部分是如何从底层TokenBlockDataset返回的三元组(source, future_target, past_target)中挑选监督目标targets只能取{self, future, past}的子集第 72-76 行 有断言源码注释第 79-87 行明确了方向约定左到右 LM 应以source为条件预测future_target右到左 LM 则预测past_targetadd_eos_for_other_targets为True时构造 self/past 目标会在source末尾补 EOS并同步调整 future/past 目标保证对齐不错位第 103-113 行双语词表支持构造时可传src_vocab与tgt_vocab两个词表_filter_vocab第 139-150 行 会把超出目标词表范围的 token 替换为目标词表的unk——这正是翻译场景中“源、目标词表可以不同”的数据层支撑add_bos_token为True时在 source 与 target 句首各补一个 BOSordered_indices第 185-193 行 用np.lexsort实现“随机序次级、长度序主级”的排序与LanguagePairDataset思路一致。其collater输出结构与平行语料略有不同单语场景下若target为空则直接复用src_tokens且填充一律在右侧第 152-173 行 文档字符串明确说明。五、Helper Datasets包装器组合出的高级数据能力文档特别指出“These datasets wrap other FairseqDataset instances and provide additional functionality”。以下逐一结合源码说明这五个文档点名的包装器。5.1BacktranslationDataset批级回译BacktranslationDataset 是合成平行数据的关键组件构造时backtranslation_fn可以传None稍后通过set_backtranslation_fn注入通常是SequenceGenerator.generate因为回译模型往往在数据加载之后才就绪关键设计__getitem__不做回译只透传目标句真正的回译发生在collater第 110-139 行——整批送进回译模型后取每条的最佳假设作为新source原句作为target再交给output_collater默认即tgt_dataset.collater拼批辅助函数 backtranslate_samples第 13-49 行 展示了输入{id, source}变为输出{id, source: 回译结果, target: 原句}的完整映射size(index)用目标句长度同时近似源、目标长度第 149-158 行因为回译前的真实源长未知。5.2ConcatDataset多语料拼接与采样比例ConcatDataset 将多个数据集首尾拼接并支持sample_ratios按倍率重复/截断各子集cumsum静态方法按int(ratio * len(e))累加出各子集的边界cumulative_sizes_get_dataset_and_sample_index用bisect.bisect_right定位全局索引落在哪个子集并用% real_size实现超界回绕从而让ratio 1的语料被重复采样ordered_indices直接按拼接后的sizes数组全局argsort让跨语料也按长度分组prefetch会按子集边界把索引切分后逐个转发并做% real_size还原supports_prefetch要求所有子集都支持预取才为真。5.3ResamplingDataset按权重逐 epoch 重采样ResamplingDataset 每个 epoch 从基础数据集重新采样weights概率权重列表默认None即均匀采样提供时会归一化后存入PlasmaArray借助 Plasma 避免多 worker 重复拷贝replaceTrue为有放回采样False为无放回size_ratioepoch 规模缩放比例必须 0。文档注释指出有放回采样时很多样本会被遗漏调高size_ratio可提升每个 epoch 对基础数据集的覆盖无放回采样时该值应小于 1batch_by_size为True时重采样后的索引会按长度重排保证“按长度分批”的特性不被打乱。5.4RoundRobinZipDatasets多任务/多语料轮询对齐RoundRobinZipDatasets 将多个数据集“拉链式”对齐到最长数据集的长度短数据集以轮询方式循环复用__len__即最长数据集的长度_map_index(key, index)对短数据集取index % len(dataset)要求构造时传入OrderedDictnum_tokens取各 key 的maxeval_key评估时只透传指定 key 的批次训练/评估两种模式共用一个数据集实例它要求先调用ordered_indices()建立各 key 的索引映射这是使用时必须注意的前置条件。5.5TransformEosDatasetEOS 的增删转换TransformEosDataset 专门处理“复用已有语料训练反方向模型”时的 EOS 不一致问题参数append_eos_to_src/remove_eos_from_src/append_eos_to_tgt/remove_eos_from_tgt且源码断言追加与移除互斥不能对同一侧同时启用转换延迟到collater中执行且首次执行时通过_check_src/_check_tgt断言“样本末尾确实不含 EOS”防止静默的数据错误构造时即预计算_src_delta/_tgt_delta让size(index)返回调整后的长度保证--max-positions过滤与批大小估计依然准确。六、Dictionary符号到连续整数的映射Dictionary 被文档定义为 “A mapping from symbols to consecutive integers”。从源码确认其核心行为特殊符号固定顺序__init__依次add_symbol注册boss、padpad、eos/s、unkunk其索引即为 0/1/2/3可经extra_special_symbols追加nspecial记录特殊符号总数index(sym)查不到时返回unk_index而非抛错第 47-51 行保证解码时 OOV 不中断string(tensor)把索引张量还原为字符串自动剔除 BOS/EOS可选去除 BPE 子词符号process_bpe_symbol、可选把unk转义为unk支持 2D 张量逐行转换add_symbol(word, n1)/update(new_dict)维护词频计数update用于把另一个词典的统计合并进来多进程分片构建词典时的标准做法finalize(threshold, nwords, padding_factor)按频率降序裁剪词表支持最小词频阈值threshold、词表上限nwords以及padding_factor默认 8将词表容量向上取整到该因子的倍数便于并行加载对齐。七、Iterators可复现、可分片、可恢复的批次迭代iterators.py 提供文档列出的四个迭代器它们共同支撑了“多 epoch 复用 多进程分片 断点续训”的数据管线。7.1CountingIterator带计数的迭代包装CountingIterator 维护count已消费元素数提供has_next()判断是否耗尽EpochBatchIterator.end_of_epoch依赖它skip(num_to_skip)用itertools.islice快速前进用于断点恢复时跳过已训练的批次take(n)截断到最多 n 个元素。7.2EpochBatchIterator相比 DataLoader 的三重增强文档中EpochBatchIterator的类说明第 136-163 行明确列出它相对torch.utils.data.DataLoader的三点优势源码逐一印证跨 epoch 复用next_epoch_itr(shuffleTrue, fix_batches_to_gpusFalse)每个 epoch 返回一个新迭代器可选在 epoch 之间洗牌可序列化state_dict()保存epoch、iterations_in_epoch、shuffle三个字段第 222-228 行load_state_dict会重建当前 epoch 的迭代器并用offset快进到中断位置实现无损续训数据分片num_shards/shard_id把批次列表切成 N 份各进程只消费自己那份。可复现性是另一个工程重点_get_iterator_for_epoch第 242-287 行 中批次洗牌的种子是self.seed epochfix_batches_to_gpusTrue时再 shard_id配合data_utils.numpy_seed上下文保证从 checkpoint 恢复后即使重放 shuffle 也能得到与中断前完全一致的批次序列。另外当数据集supports_prefetch时洗牌与分片完成后会立即调用self.dataset.prefetch([...])预取整个 epoch 所需数据num_workers 0时最终仍由DataLoader承载多进程加载。7.3GroupedIterator与ShardedIteratorGroupedIterator第 290-318 行把迭代器按chunk_size分块输出长度取ceil(len/chunk_size)末块不足chunk_size时照常返回仅空块触发StopIterationShardedIterator第 321-353 行用itertools.islice(iterable, shard_id, None, num_shards)实现交错分片第 i 个分片取第 i、iN、i2N… 个元素而非连续切块并用zip_longestfill_value保证各分片长度一致向上取整避免某个进程提前结束。八、小结数据管线的调用链与延伸阅读综合以上源码InfoXLM 内置 fairseq 的数据加载链路可以概括为原始语料 → (preprocess 二进制化) → IndexedDataset → LanguagePairDataset / MonolingualDataset定义 __getitem__/size/num_tokens → 可选 Helper 包装Concat / Resampling / RoundRobinZip / TransformEos / Backtranslation → ordered_indices() 决定批内/批间顺序 → EpochBatchIteratorseedepoch 洗牌、num_shards 分片、state_dict 恢复 → collater 拼出 net_input / target → 训练器 forward文档 data.rst 中每个 autoclass 条目都对应上文给出的真实源码文件可继续深入的方向包括预取与索引存储的具体实现indexed_dataset.pyMMapIndexedDataset等prefetch能力的提供方批处理底层工具data_utils.pycollate_tokens的填充与 EOS 移位逻辑;数据二进制的生成入口preprocess.py 与 binarizer.pyInfoXLM 模型侧如何消费这些数据集src-infoxlm 目录下的任务定义。需要强调的适用前提以上描述均基于本仓库infoxlm/fairseq目录下的 fairseq 源码快照MIT 协议LICENSE参数默认值与行为以该版本源码为准仓库根目录的 README 亦说明 unilm 仓库聚合了多个模型子项目本文仅覆盖其中 InfoXLM 所携带的 fairseq 数据模块。【免费下载链接】unilmLarge-scale Self-supervised Pre-training Across Tasks, Languages, and Modalities项目地址: https://gitcode.com/GitHub_Trending/un/unilm创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考