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 TerabyteTB点击日志数据集的完整预处理流程如何用 Apache Beam 完成数据分片再平衡、词表生成与特征变换最终产出可直接被 DLRM / DCN v2 排序模型训练的 TSV 数据。读完后你可以复制文中的命令在 Google Cloud Dataflow 或本地 DirectRunner 上完成全流程预处理并理解每步背后的 Beam/TFT 实现细节。流程总览原始 Criteo TB 数据集由 24 天day_0.gz~day_23.gz约 1TB 的压缩点击日志组成无法直接用于训练仓库中的排序模型。预处理脚本说明 给出的流水线共四步对应两个核心脚本将原始数据下载并上传到 GCS或本地路径用 shard_rebalancer.py 把原始 train/test 数据切分为固定数量的分片均衡后续并行处理负载用 criteo_preprocess.py 分两趟生成词汇表并变换特征产出训练/测试 TSV可选再次用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 项目 IDexport STORAGE_BUCKETgs://bucket-name export PROJECTmy-gcp-project export REGIONus-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 定义参数默认值说明--input_path必填输入文件支持通配符会读取该路径下所有分片--output_path必填输出路径前缀输出文件名由 Beam 按 shard 数自动编号--num_output_files256输出分片数量训练集用 1024、测试集用 64--filetypetfrecord文件类型取值{tfrecord, csv}本流程用csv--runnerDirectRunnerBeam Runner可选DirectRunner/DataflowRunner--project/--region无Dataflow 项目 ID 与区域DirectRunner 下被忽略从源码结构看shard_rebalancer.py的核心逻辑非常薄CSV 模式就是ReadFromText(input_path)后接WriteToText(output_path, num_shardsnum_output_files)TFRecord 模式则换成ReadFromTFRecordWriteToTFRecord使用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_modeTrue只生成词表文件第二趟不写词表、只把词表应用到数据上。两趟必须使用相同的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 参数全集参数默认值说明--input_path必填输入路径需覆盖全部数据以保证稀疏词表完整--output_path必填输出路径同时被用作 Dataflow 的 staging/temp 前缀--temp_dir必填临时元数据目录词表字典存于此两趟运行必须一致--csv_delimeter\t输入/输出字段分隔符与下游 TSV 读取保持一致--vocab_gen_modeFalse置位后扫全量数据、不写 CSV 输出仅产出词表--runnerDirectRunnerBeam Runner可选DirectRunner/DataflowRunner--project/--region无Dataflow 项目与区域DirectRunner 下忽略--max_vocab_size10,000,000哈希桶上限类别特征转整数后对其取模词表生成的源码实现从源码看词表由 TensorFlow TransformTFT在 Beam 上计算。compute_vocab_fn对每个类别特征调用tft.compute_and_apply_vocabulary(xinputs[key], vocab_filenamefeature_{idx}_vocab)依赖 TFT 的隐式行为把词表写到temp_dir下tftransform_tmp/feature_??_vocab随后在tft_beam.AnalyzeAndTransformDataset这一步统一完成统计与分析。脚本先建好tft_beam.Context(temp_dirargs.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([], dtypetf.int64)数值特征与标签为tf.float32的FixedLenFeature——因为在此之前数据已经过纯 Beam 阶段的清洗见下文TFT 看到的是规整的定长数值可以安全地使用 FixedLenFeature schema。第三步两趟之外的清洗逻辑纯 Beam 阶段在把数据交给 TFT 之前流水线先用三个轻量beam.DoFn做文本级清洗这是理解输入/输出格式的关键FillMissingCriteo 原始行用\t分隔空值表示缺失。该算子把每个空字段替换为字符串0保证每行字段数恒定NegsToZeroLog对第 1 ~ 13 个数值特征下标i 0 and i 13先把负值置零再取np.log(int(val) 1)即log(x1)归一化输出为浮点字符串HexToIntModRange对第 14 列起的类别特征下标i 13把十六进制字符串按int(val, 16)解析为十进制再对max_vocab_size取模得到落在[0, max_vocab_size)的整数。这三步与上层文档 ranking/README.md 中总结的处理规则完全一致整型特征依次做缺失置零 → 负值置零 → log(x1)类别特征做十六进制转整数 → 可选地按模数降维feature_value % MAX_INDEX后者用于压缩超大嵌入表的规模。清洗后的行交给tfxio.BeamRecordCsvTFXIOphysical_formattext按固定列序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(x1)、类别分桶的预处理。一个容易踩坑的细节第二趟 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 启动 DLRMinteraction: dot或 DCN v2interaction: cross训练。本地运行与排错要点Runner 选择--runner DirectRunner时无需--project/--regioncriteo_preprocess.py在 DirectRunner 下会用direct_num_workersos.cpu_count()与multi_threading模式多线程处理。小样本验证建议先按 README 提示把max_vocab_size调小、输入路径指向本地目录。两趟一致性词表模式的input_path必须覆盖 traintest第二趟的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(x1)的浮点数——这正是CriteoTsvReader所假设的输入形态。小结official/recommendation/ranking/preprocessing提供了一条可复现的 Criteo TB 数据处理链路shard_rebalancer.py负责输入输出的分片均衡criteo_preprocess.py用 Beam 文本清洗 TFT 词表 的组合完成缺失填充、log(x1)归一化、十六进制转整数取模和词表映射两趟作业。产出的 TSV 与vocab_sizes配置对接后即可直接喂给仓库中的 DLRM / DCN v2 排序模型训练入口。【免费下载链接】modelsModels and examples built with TensorFlow项目地址: https://gitcode.com/GitHub_Trending/mode/models创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考