TensorFlow Models Ranking 模块：基于 Apache Beam 的 Criteo TB 数据集分布式预处理指南

TensorFlow Models Ranking 模块：基于 Apache Beam 的 Criteo TB 数据集分布式预处理指南

【免费下载链接】modelsModels and examples built with TensorFlow项目地址: https://gitcode.com/GitHub_Trending/mode/models

本篇介绍 TensorFlow Models 仓库official/recommendation/ranking中 Criteo Terabyte（TB）点击日志数据集的完整预处理流程：如何用 Apache Beam 完成数据分片再平衡、词表生成与特征变换，最终产出可直接被 DLRM / DCN v2 排序模型训练的 TSV 数据。读完后你可以复制文中的命令在 Google Cloud Dataflow 或本地 DirectRunner 上完成全流程预处理，并理解每步背后的 Beam/TFT 实现细节。

流程总览

原始 Criteo TB 数据集由 24 天（day_0.gz~day_23.gz）约 1TB 的压缩点击日志组成，无法直接用于训练仓库中的排序模型。预处理脚本说明 给出的流水线共四步，对应两个核心脚本：

1. 将原始数据下载并上传到 GCS（或本地路径）；
2. 用 shard_rebalancer.py 把原始 train/test 数据切分为固定数量的分片，均衡后续并行处理负载；
3. 用 criteo_preprocess.py 分两趟生成词汇表并变换特征，产出训练/测试 TSV；
4. （可选）再次用shard_rebalancer.py对产出数据做再分片，得到最终训练桶。

整个流程基于 Apache Beam 构建，既可以在本地以 DirectRunner 运行（前提是主机有足够 CPU/内存/存储），也可以通过DataflowRunner提交到 Google Cloud Dataflow 分布式执行，这也是官方推荐 Criteo TB 这种规模数据的方式。

环境准备与依赖安装

安装预处理依赖

在official/recommendation/ranking/preprocessing目录下执行：

python3 setup.py install

这里的 setup.py 定义了一个名为criteo_preprocessing的包，install_requires中声明了核心依赖tensorflow-transform。它同时承担 Dataflow 提交时的setup_file角色——criteo_preprocess.py在构建 Dataflow 流水线选项时会把./setup.py作为setup_file传入，这样 Dataflow worker 节点上也能安装tensorflow-transform，而 shard_rebalancer.py 则只依赖 Apache Beam 与 TensorFlow 本体。

设置环境变量

将bucket-name替换为你的 Cloud Storage 桶名，my-gcp-project替换为你的 GCP 项目 ID：

export STORAGE_BUCKET=gs://bucket-name export PROJECT=my-gcp-project export REGION=us-central1

注意：若要在本地运行，上述环境变量并不需要，gs://bucket-name可以直接换成本地路径；本地跑时建议把max_vocab_size传一个更小的值以控制内存占用（本地单机处理全量 TB 级数据的词表生成不现实）。

组织原始数据

从 Criteo 官方下载原始数据后，按以下结构组织：

• day_0.gz, day_1.gz, ..., day_22.gz放在${STORAGE_BUCKET}/criteo_raw/train/
• day_23.gz放在${STORAGE_BUCKET}/criteo_raw/test/

即前 23 天训练、第 24 天评估，这与上层 ranking/README.md 中"days 1-23 for training and day 24 for evaluation" 的划分一致。

第一步：Shard 分片再平衡

python3 shard_rebalancer.py \ --input_path "${STORAGE_BUCKET}/criteo_raw/train/*" \ --output_path "${STORAGE_BUCKET}/criteo_raw_sharded/train/train" \ --num_output_files 1024 --filetype csv --runner DataflowRunner \ --project ${PROJECT} --region ${REGION}

python3 shard_rebalancer.py \ --input_path "${STORAGE_BUCKET}/criteo_raw/test/*" \ --output_path "${STORAGE_BUCKET}/criteo_raw_sharded/test/test" \ --num_output_files 64 --filetype csv --runner DataflowRunner \ --project ${PROJECT} --region ${REGION}

参数说明（来自脚本内 argparse 定义）：

从源码结构看，shard_rebalancer.py的核心逻辑非常薄：CSV 模式就是ReadFromText(input_path)后接WriteToText(output_path, num_shards=num_output_files)；TFRecord 模式则换成ReadFromTFRecord+WriteToTFRecord（使用ProtoCoder(tf.train.Example)）。也就是说它本质上是一个"读入所有行—随机写出到 N 个新文件"的再分片器，作用是把 Criteo 按天组织的 24 个大文件打散成大量均衡小分片，避免后续 Beam 流水线中出现单个 day 文件成为处理热点。作业名会自动生成为shard-rebalancer-<时间戳>，Dataflow 的staging_location和temp_location都派生自output_path/tmp/下。

第二步：生成词汇表（vocab_gen_mode）

预处理的核心是 criteo_preprocess.py，其文件头注释明确说明：对原始 Criteo 数据必须运行两趟——第一趟vocab_gen_mode=True只生成词表文件，第二趟不写词表、只把词表应用到数据上。两趟必须使用相同的temp_dir。

第一趟命令：

python3 criteo_preprocess.py \ --input_path "${STORAGE_BUCKET}/criteo_raw_sharded/*/*" \ --output_path "${STORAGE_BUCKET}/criteo/" \ --temp_dir "${STORAGE_BUCKET}/criteo_vocab/" \ --vocab_gen_mode --runner DataflowRunner --max_vocab_size 5000000 \ --project ${PROJECT} --region ${REGION}

注意--input_path覆盖了 train 和 test 所有分片（criteo_raw_sharded/*/*），脚本参数注释也强调：input_path必须覆盖全部数据（train、test、eval 都要），否则稀疏特征词表不完整。词表结果落在${STORAGE_BUCKET}/criteo_vocab/tftransform_tmp/feature_??_vocab文件中，词表大小可以用wc -l <feature_vocab_file>查看。

criteo_preprocess.py 参数全集

词表生成的源码实现

从源码看，词表由 TensorFlow Transform（TFT）在 Beam 上计算。compute_vocab_fn对每个类别特征调用tft.compute_and_apply_vocabulary(x=inputs[key], vocab_filename="feature_{idx}_vocab")，依赖 TFT 的隐式行为把词表写到temp_dir下（tftransform_tmp/feature_??_vocab），随后在tft_beam.AnalyzeAndTransformDataset这一步统一完成统计与分析。脚本先建好tft_beam.Context(temp_dir=args.temp_dir)上下文，再进入 Beam 流水线。

特征命名约定也在源码中固定下来：Criteo TB 的 39 个特征里，第 1 列是标签clicked，随后 13 列是数值特征int-feature-1~int-feature-13，其余 26 列是类别特征categorical-feature-14~categorical-feature-39。这些列最终会被声明为 TFT 的输入 schema：类别特征为tf.io.FixedLenFeature([], dtype=tf.int64)，数值特征与标签为tf.float32的FixedLenFeature——因为在此之前数据已经过纯 Beam 阶段的清洗（见下文），TFT 看到的是规整的定长数值，可以安全地使用 FixedLenFeature schema。

第三步：两趟之外的清洗逻辑（纯 Beam 阶段）

在把数据交给 TFT 之前，流水线先用三个轻量beam.DoFn做文本级清洗，这是理解输入/输出格式的关键：

1. FillMissing：Criteo 原始行用\t分隔，空值表示缺失。该算子把每个空字段替换为字符串"0"，保证每行字段数恒定；
2. NegsToZeroLog：对第 1 ~ 13 个数值特征（下标i > 0 and i <= 13），先把负值置零，再取np.log(int(val) + 1)，即log(x+1)归一化，输出为浮点字符串；
3. HexToIntModRange：对第 14 列起的类别特征（下标i > 13），把十六进制字符串按int(val, 16)解析为十进制，再对max_vocab_size取模，得到落在[0, max_vocab_size)的整数。

这三步与上层文档 ranking/README.md 中总结的处理规则完全一致：整型特征依次做"缺失置零 → 负值置零 → log(x+1)"；类别特征做"十六进制转整数 → 可选地按模数降维"（feature_value % MAX_INDEX），后者用于压缩超大嵌入表的规模。

清洗后的行交给tfxio.BeamRecordCsvTFXIO（physical_format="text"）按固定列序clicked, int-feature-1..13, categorical-feature-14..39解析成张量，然后进入tft_beam.AnalyzeAndTransformDataset。第二趟（非词表模式）中，转换函数是apply_vocab_fn：对每个类别特征调用tft.apply_vocabulary(inputs[key], vocab_fn)，其中vocab_fn就是第一趟落盘的feature_??_vocab文件路径；数值特征与标签则原样透传。

python3 criteo_preprocess.py \ --input_path "${STORAGE_BUCKET}/criteo_raw_sharded/train/*" \ --output_path "${STORAGE_BUCKET}/criteo/train/train" \ --temp_dir "${STORAGE_BUCKET}/criteo_vocab/" \ --runner DataflowRunner --max_vocab_size 5000000 \ --project ${PROJECT} --region ${REGION}

python3 criteo_preprocess.py \ --input_path "${STORAGE_BUCKET}/criteo_raw_sharded/test/*" \ --output_path "${STORAGE_BUCKET}/criteo/test/test" \ --temp_dir "${STORAGE_BUCKET}/criteo_vocab/" \ --runner DataflowRunner --max_vocab_size 5000000 \ --project ${PROJECT} --region ${REGION}

两趟变换完成后，若不在vocab_gen_mode，脚本用tft.coders.CsvCoder按相同列序与\t分隔符把张量编码回 CSV 并WriteToText到output_path。最终产出就是 TSV 文件：每行依次为标签、13 个 log 归一化后的浮点数值特征、26 个整型类别特征，这恰好是下游 data_pipeline.py 中CriteoTsvReader的解析假设——它用tf.io.decode_csv(..., field_delim='\t', na_value='-1')按label + dense + sparse的顺序解码，并注释写明输入必须已经过"缺失置零、负值置零、log(x+1)、类别分桶"的预处理。

一个容易踩坑的细节：第二趟 train/test 预处理时max_vocab_size必须与第一趟生成词表时使用的值一致（示例统一为 5,000,000），因为取模范围与词表统计范围耦合。

第四步（可选）：再平衡产出数据

python3 shard_rebalancer.py \ --input_path "${STORAGE_BUCKET}/criteo/train/*" \ --output_path "${STORAGE_BUCKET}/criteo_balanced/train/train" \ --num_output_files 8192 --filetype csv --runner DataflowRunner \ --project ${PROJECT} --region ${REGION}

python3 shard_rebalancer.py \ --input_path "${STORAGE_BUCKET}/criteo/test/*" \ --output_path "${STORAGE_BUCKET}/criteo_balanced/test/test" \ --num_output_files 1024 --filetype csv --runner DataflowRunner \ --project ${PROJECT} --region ${REGION}

最终训练/测试数据落在：

• ${STORAGE_BUCKET}/criteo_balanced/train/
• ${STORAGE_BUCKET}/criteo_balanced/test/

此之后其他中间桶（criteo_raw*、criteo*等）均可删除。这一步把预处理输出的分片重新打散成 8192 / 1024 个更均匀的分片，让训练时的文件级分片调度更平滑。

与训练配置的衔接：vocab_sizes 从哪来

预处理产出的词表大小直接决定训练侧的model.vocab_sizes。仓库中的 configs/config.py 给出了 Criteo TB 的标准取值：

vocab_sizes = [ 39884406, 39043, 17289, 7420, 20263, 3, 7120, 1543, 63, 38532951, 2953546, 403346, 10, 2208, 11938, 155, 4, 976, 14, 39979771, 25641295, 39664984, 585935, 12972, 108, 36, ]

这 26 个数与 26 个类别特征一一对应（顺序与输入数据列序一致），并用于DataConfig/ModelConfig；其中NUM_TRAIN_EXAMPLES = 4,195,197,692、NUM_EVAL_EXAMPLES = 89,137,318也定义了默认的train_steps、validation_steps。dlrm_criteo_tpu.yaml 与 dcn_v2_criteo_tpu.yaml 两份 YAML 配置里同样内嵌了这组vocab_sizes，只需把train_data.input_path/validation_data.input_path指向上一步产出的${STORAGE_BUCKET}/criteo_balanced/train/*与.../test/*，即可通过 train.py 启动 DLRM（interaction: 'dot'）或 DCN v2（interaction: 'cross'）训练。

本地运行与排错要点

• Runner 选择：--runner DirectRunner时无需--project/--region；criteo_preprocess.py在 DirectRunner 下会用direct_num_workers=os.cpu_count()与multi_threading模式多线程处理。小样本验证建议先按 README 提示把max_vocab_size调小、输入路径指向本地目录。
• 两趟一致性：词表模式的input_path必须覆盖 train+test；第二趟的temp_dir必须指向第一趟写入词表的同一目录，否则apply_vocabulary找不到feature_??_vocab文件。
• Dataflow 临时目录：criteo_preprocess.py会把staging_location/temp_location放在output_path/tmp/下，作业名形如criteo-preprocessing-<时间戳>，可在 GCP 控制台按此检索作业日志；setup_file=./setup.py保证 worker 上装好tensorflow-transform。
• 输出格式核对：抽样检查产出的 TSV 行数/列数应为 39 列（1 标签 + 13 数值 + 26 类别），类别值落在[0, max_vocab_size)，数值列应已是对数尺度（log(x+1)）的浮点数——这正是CriteoTsvReader所假设的输入形态。

小结

official/recommendation/ranking/preprocessing提供了一条可复现的 Criteo TB 数据处理链路：shard_rebalancer.py负责输入输出的分片均衡，criteo_preprocess.py用 "Beam 文本清洗 + TFT 词表" 的组合完成缺失填充、log(x+1)归一化、十六进制转整数取模和词表映射两趟作业。产出的 TSV 与vocab_sizes配置对接后，即可直接喂给仓库中的 DLRM / DCN v2 排序模型训练入口。

【免费下载链接】modelsModels and examples built with TensorFlow项目地址: https://gitcode.com/GitHub_Trending/mode/models