tensorflow/models NLP 训练数据怎么在离线预处理脚本与 TFDS 在线读取之间选型？

在 tensorflow/models 仓库的 official/nlp（TF-NLP）里训练 BERT 等预训练模型时，输入数据有两条被官方文档支持的供给方式：一是用 Python 预处理脚本把原始语料离线 tokenize、构造 tf.Example 再落盘，训练时通过 DataConfig.input_path 读文件；二是直接指定 TFDS 数据集，由 input_reader.py 在 tf.data 输入管线内用 TF.Text 完成 tokenize 和预处理。数据加工文档把这两种方式描述为“flexible ways to provide training data to the input pipeline”，并且二者最终都收敛到同一个 DataConfig 对象上，互斥使用。本文给出两条路径的完整操作步骤、可直接复用的命令，以及依据仓库文档可核对的选型判断点。

两条路径的共同落点：DataConfig

无论选哪条路径，数据配置都写在实验配置（YAML 或 Python factory）的 task.train_data / task.validation_data 里，类型是 config_definitions.py 中的 DataConfig。与选型直接相关的字段：

• input_path：离线预处理产出的文件路径。可以是单个路径/通配符、逗号分隔的多路径（"a, b, c"）、列表，或上述三者的字典（用于多数据集混合）。
• tfds_name：TFDS 数据集名，指定了 tfds_name 时不能同时指定 input_path。
• tfds_split：从 TFDS 加载哪个 split，指定 tfds_name 时为必填。
• tfds_data_dir：TFDS 数据读写的目录。
• tfds_as_supervised、tfds_skip_decoding_feature：TFDS 专属的加载行为开关。

InputReader 在初始化时强制这条互斥规则（input_reader.py）：

if params.input_path and params.tfds_name:
  raise ValueError('At most one of `input_path` and `tfds_name` can be '
                   'specified, but got %s and %s.' % ...)

也就是说“同时写 input_path 和 tfds_name”会被直接拒绝，选型必须在配置阶段二选一。

路径一：离线预处理脚本 + input_path 读 tfrecord

适合语料来自自定义来源、或想沿用原 BERT 论文 Wiki + Books 配方的场景。pretrain.md给出的完整流程是：

第 1 步：跑预处理脚本生成训练文件。 脚本是 create_pretraining_data.py，文档中的示例命令（$WORKING_DIR、$BERT_DIR 为文档示例中的环境变量，需替换为你自己的磁盘位置）：

export WORKING_DIR='local disk or cloud location'
export BERT_DIR='local disk or cloud location'
python models/official/nlp/data/create_pretraining_data.py \
  --input_file=$WORKING_DIR/input/input.txt \
  --output_file=$WORKING_DIR/output/tf_examples.tfrecord \
  --vocab_file=$BERT_DIR/wwm_uncased_L-24_H-1024_A-16/vocab.txt \
  --do_lower_case=True \
  --max_seq_length=512 \
  --max_predictions_per_seq=76 \
  --masked_lm_prob=0.15 \
  --random_seed=12345 \
  --dupe_factor=5

关键参数（均来自脚本内 absl flag 定义）：

• --input_file / --output_file：两者都是必填项（mark_flag_as_required）；
• --max_seq_length：文档特别强调它必须与你后续预训练时指定的序列长度参数一致；
• --tokenization：WordPiece（Canonical BERT）或 SentencePiece（ALBERT），配套 --vocab_file 或 --sp_model_file；
• --do_whole_word_mask、--max_ngram_size：n-gram 掩码要求同时设置 --do_whole_word_mask=True；
• --masked_lm_prob（默认 0.15）、--short_seq_prob（默认 0.1）、--dupe_factor（默认 10）、--gzip_compress、--use_v2_feature_names。

微调数据另有 create_finetuning_data.py，用于分类/检索/序列标注等任务的离线构造。

第 2 步：把产出路径写进实验配置。 以 wiki_books_pretrain.yaml 为例，train_data.input_path 形如：

input_path: '[Your processed wiki data path]*,[Your processed books data path]*'

方括号内容是配置里预留的占位符，需替换为你第 1 步的实际输出路径前缀；文档说明当数据有多个 shard 时可以用 * 通配符把多个文件一次纳入。同时把 seq_length 等掩码相关超参与预处理参数对齐（该 yaml 中 seq_length: 512、max_predictions_per_seq: 76，与上面的脚本参数一致）。

第 3 步：启动训练。 运行 bert/pretraining 实验（pretraining_experiments.py），文档示例命令：

export OUTPUT_DIR=gs://some_bucket/my_output_dir
export PARAMS=$PARAMS,runtime.distribution_strategy=tpu

python3 train.py \
 --experiment=bert/pretraining \
 --mode=train_and_eval \
 --model_dir=$OUTPUT_DIR \
 --config_file=configs/models/bert_en_uncased_base.yaml \
 --config_file=configs/experiments/wiki_books_pretrain.yaml \
 --tpu=${TPU_NAME} \
 --params_override=$PARAMS

pretrain.md中的 gs://some_bucket/my_output_dir、${TPU_NAME} 是文档示例值，替换为你自己的输出目录与 TPU 实例名。train.py 对 experiment、mode、model_dir 三个 flag 做了 mark_flags_as_required，缺了任何一个都会直接报错。

路径二：TFDS 在线读取 + TF.Text 管线内预处理

input_reader.py 内置了 TFDS 读取路径（_read_tfds），配合 wiki_tfds_pretrain.yaml，数据在 tf.data 管线内直接 tokenize，不需要先落盘。该 yaml 的关键字段：

tfds_name: wikipedia/20201201.en
tfds_split: train
vocab_file_path: 'Please provide the vocab file path.'

vocab_file_path 是明确留给读者填写的占位符，替换为你模型的词表路径。这条路径在 pretraining_experiments.py 中对应 bert/text_wiki_pretraining 实验，内部使用 pretrain_text_dataloader.BertPretrainTextDataConfig，其数据管线（pretrain_text_dataloader.py）用 tf_text.BertTokenizer、tf_text.WaterfallTrimmer、tf_text.RegexSplitter、tf_text.mask_language_model 等算子在 GPU/TPU 上完成切分、拼接与掩码，不依赖本地预处理产物。

文档示例的启动命令（$OUTPUT_DIR、$BERT_DIR、${TPU_NAME} 均为文档示例变量，替换后使用）：

export OUTPUT_DIR=gs://some_bucket/my_output_dir
export BERT_DIR=~/cased_L-12_H-768_A-12
export PARAMS=$PARAMS,task.validation_data.vocab_file_path=$BERT_DIR/vocab.txt
export PARAMS=$PARAMS,task.train_data.vocab_file_path=$BERT_DIR/vocab.txt
export PARAMS=$PARAMS,runtime.distribution_strategy=tpu

python3 train.py \
 --experiment=bert/text_wiki_pretraining \
 --mode=train_and_eval \
 --model_dir=$OUTPUT_DIR \
 --config_file=configs/experiments/wiki_tfds_pretrain.yaml \
 --tpu=${TPU_NAME} \
 --params_override=$PARAMS

注意 pretrain.md 明确标注该示例“only wikipedia english corpus is used”，即官方 TFDS 示例只覆盖维基百科英文语料。

选型判断点

文档没有给出通用“优劣”结论，但提供了以下可核对的判断依据，按你的实际情况对号入座：

判断维度	离线预处理脚本路径	TFDS 在线读取路径
语料来源	自定义原始文本（Wiki dump、BookCorpus 或自有语料），可组合多来源（如 input_path 字典做命名混合）	语料已是 TFDS 数据集（文档示例为 wikipedia/20201201.en）
tokenize 时机	离线完成，产出 tf.Example proto 文件（tfrecord）	在线完成，TF.Text 算子在 tf.data 管线内 tokenize
掩码/预处理控制参数	masked_lm_prob、short_seq_prob、dupe_factor、whole-word/n-gram 掩码等在脚本 flag 中控制	由 BertPretrainTextDataConfig 及 yaml 字段控制，如 use_whole_word_masking: true
与 TPU 动态序列配合	pretrain_dynamic_dataloader.py 面向 tokenized 数据集，配合 --enable_tf_data_service 与全局 bucketizing；文档报告在典型文本上相对静态定长输入可获得 50%–90% 的训练加速	文档的 TFDS 示例未包含动态序列路径
数据读取细节	文件级 shuffle/shard，按 input_path 通配符匹配	tfds_name 以 mldataset. 开头时走 tfds.load，否则走 tfds.builder；当 split 的 shard 数小于输入管线数时先整读再在 host memory 里 shard

归纳成两条可执行规则：

1. 语料不是 TFDS 数据集、或需要自定义多源混合与掩码配方 → 走离线脚本，把产出写进 input_path；
2. 语料在 TFDS 中、希望省去本地预处理步骤（文档称 TFDS 路径是为 “convenience and consolidation” 提供统一读取）→ 走 tfds_name/tfds_split，用 TF.Text 在管线内处理。

另外，data_processing.md 指出 TPU 训练时整个 DataLoader 的 load 方法会运行在 TPU worker 上，函数内不能访问 task 属性等外部资源；两条路径的 DataLoader 都要满足这条约束。处理原始文本特征时，可参考 sentence_prediction_dataloader.py（GLUE 微调 + TFDS 原始文本）和 pretrain_dynamic_dataloader.py（tokenized 数据集 + tf.data service）作为实现范本。

验证与错误判定

以下报错和日志都来自仓库源码，可用于逐步确认配置是否生效：

1. 预处理产出确认：create_pretraining_data.py 在写文件前会输出日志 *** Writing to output files *** 并逐个打印输出路径（脚本 main 函数）。看到输出文件列表、且磁盘/存储中出现对应 tfrecord 文件，说明第 1 步完成。
2. 路径匹配检查：input_path 里的通配符若一个文件都匹配不到，match_files 会抛出 ValueError: '<pattern>' does not match any files.（input_reader.py）。出现该错误说明预处理产物路径写错或 shard 命名对不上，先回到第 1 步核对输出路径。
3. TFDS 配置完整性：指定了 tfds_name 却没指定 tfds_split 时，InputReader 抛出 ValueError: '`tfds_name` is ..., but `tfds_split` is not specified.'；同时指定 input_path 与 tfds_name 时抛出 At most one of input_path and tfds_name can be specified。sentence_prediction_dataloader.py 中还有同样的成对校验逻辑（tfds_name 与 tfds_split 必须同时给出、与 input_path 互斥）。
4. 训练入口检查：train.py 必填 --experiment、--mode、--model_dir；TPU 场景通过 --params_override 传入 runtime.distribution_strategy=tpu。训练能否正常进入 loop 即为该路径配置成立的最终信号。

限制与边界

• 文档示例中的 TFDS 预训练仅覆盖维基百科英文语料；换成其他语料时，tfds_name/tfds_split 是否可用取决于 TFDS 侧是否有对应数据集，仓库文档未提供更多现成示例。
• tfds_data_dir 是 TFDS 数据“读/写”目录（DataConfig 字段说明），首次从 TFDS 拉取数据时会占用该目录空间，选择时注意存储位置。
• 文件数少于输入管线数时，input_reader.py 会警告 The number of files ... is less than the number of input pipelines ... Please consider sharding your data into more files.（input_reader.py），提示把离线预处理产物切分更多 shard。
• 共享 tf.data service 等高级特性（enable_shared_tf_data_service_between_parallel_trainers 等）属于跨并行 trainer 调参场景，与本场景的基础选型无关，这里不展开。