TensorFlow Model Garden Ranking 指南：DLRM 与 DCN v2 排序模型的数据预处理、配置系统与 TPU/GPU 训练实战

TensorFlow Model Garden Ranking 指南DLRM 与 DCN v2 排序模型的数据预处理、配置系统与 TPU/GPU 训练实战【免费下载链接】modelsModels and examples built with TensorFlow项目地址: https://gitcode.com/GitHub_Trending/mode/models本文基于official/recommendation/ranking/README.md展开系统讲解 Model Garden 中排序Ranking模型的完整落地链路如何准备 Criteo 级数据集并完成特征预处理如何通过params_override以 YAML 方式驱动 DLRM点积特征交互与 DCN v2交叉网络特征交互的训练与评估并结合仓库源码剖析双优化器、学习率调度、TPU Embedding 分层策略与数据分片管道的实现细节。读完后你将能够独立完成从原始点击日志到可训练 TSV 数据、再到 GPU/TPU 上可复现实验的全流程。1. 排序模型总览DLRM 与 DCN v2该模块是 DLRM 与 DCN v2 两种深度排序模型的实现面向 CTR点击率预估这类任务。其输入由数值特征numerical features和类别特征categorical features组成输出为标量例如点击概率。三类硬件的定位差异是理解整个设计的关键CPU适合大规模稀疏 Embedding 查找lookupGPU适合深度网络MLP的快速计算TPU两者兼顾。在 TPU 上训练时类别特征使用 TensorFlow Recommenders 库中的TPUEmbedding 层。TPU Embedding 支持大表快速查找表大小随 TPU Pod 规模线性扩展v3-8 约 90 GB、v3-512 约 5.6 TB、v3-2048 Pod 可达 22.4 TB。需要注意的分工是模型本体代码位于 TensorFlow Recommenders 库的tensorflow_recommenders/experimental/models中而本仓库承担输入管道、配置体系与训练循环。入口文件为 train.py。2. 环境准备2.1 获取代码git clone https://github.com/tensorflow/models.git export PYTHONPATH$PYTHONPATH:$(pwd)/models也可以使用预装好的 Google Cloud VM 镜像。2.2 安装 TensorFlow Recommenders必须安装 TensorFlow Recommenders 库安装方式取决于所用 TensorFlow 版本tf-nightly 用户由于 night 版与 Recommenders 的依赖锁定冲突需通过--no-deps跳过依赖安装pip install tensorflow-recommenders --no-deps稳定版 TensorFlow 2.4直接安装即可pip install tensorflow-recommenders3. 数据集准备3.1 可选数据集与合成数据模型可在多种数据集上训练常用的是Criteo Terabyte与Criteo Kaggle两个公开点击日志数据集。此外可以完全不用真实数据——设置use_synthetic_data: true即可基于随机张量生成合成数据用于快速打通整条训练链路合成数据管道见 data_pipeline.py 中CriteoTsvReader._generate_synthetic_data按vocab_sizes生成随机稀疏 ID 并构造 0/1 标签。真实数据集下载需按 Criteo 官方说明进行。注意 Criteo Terabyte 体量巨大约 1TB。3.2 特征预处理步骤详细流程见 preprocessing/README.md预处理脚本基于 Apache Beam可在 Google Cloud Dataflow 上分布式运行也可本地以 DirectRunner 运行前提是本机有足够 CPU/内存/存储。核心规则如下整数数值特征按顺序执行缺失值替换为 0负值替换为 0对特征做log(x1)变换结果为tf.float32。类别特征分桶bucketize为tf.int32可选地将整数进一步哈希到更低维度以降低大表的体积。简单的取模哈希即可例如feature_value % MAX_INDEX。预处理得到的词表大小vocabulary sizes会通过model.vocab_sizes配置传入模型训练器。注意README 示例中给出的具体取值仅对 Criteo Terabyte 数据集有效。完整数据集由 24 个日期目录组成可将其划分为训练/验证集例如第 123 天用于训练、第 24 天用于评估。预处理流水线本身包含四步完整命令见 preprocessing/README.md将原始day_*.gz文件组织到 GCS 桶的 train/test 目录用 shard_rebalancer.py 把原始数据分片训练 1024 个文件、测试 64 个文件用 criteo_preprocess.py 先生成各特征词表--vocab_gen_mode再对 train/test 数据做编码转换可选再次 re-balance 数据集。依赖安装为python3 setup.py install其中 setup 文件位于 preprocessing/setup.py。3.3 训练/验证数据的 TSV 格式训练与验证数据须保存为大量 TSV 文件格式为数值特征、类别特征、标签。每行第一列是标签0 或 1随后num_dense_features个数值特征再后是vocab_sizes个类别特征第 i 个类别特征必须是[0, vocab_sizes[i])区间内的整数。数据目录约定为两个子目录${DATA_DIR}/train与${DATA_DIR}/eval。4. 训练与评估两种模型的差异仅体现在interaction配置上训练DLRM使用点积特征交互interaction: dot训练DCN v2使用交叉层interaction: cross此外源码还支持interaction: multi_layer_dcn多层低秩 DCN对应 DLRM DCN v2 组合配置见 task.py 中build_model对三种交互的分支处理。4.1 TPU 训练完整命令export TPU_NAMEmy-dlrm-tpu export EXPERIMENT_NAMEmy_experiment_name export BUCKET_NAMEgs://my_dlrm_bucket export DATA_DIR${BUCKET_NAME}/data export EMBEDDING_DIM32 python3 models/official/recommendation/ranking/train.py --modetrain_and_eval \ --model_dir${BUCKET_NAME}/model_dirs/${EXPERIMENT_NAME} --params_override runtime: distribution_strategy: tpu task: use_synthetic_data: false train_data: input_path: ${DATA_DIR}/train/* global_batch_size: 16384 validation_data: input_path: ${DATA_DIR}/eval/* global_batch_size: 16384 model: num_dense_features: 13 bottom_mlp: [512,256,${EMBEDDING_DIM}] embedding_dim: ${EMBEDDING_DIM} top_mlp: [1024,1024,512,256,1] interaction: dot 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] trainer: use_orbit: true validation_interval: 85352 checkpoint_interval: 85352 validation_steps: 5440 train_steps: 256054 steps_per_loop: 1000 要点说明--mode支持train、eval、train_and_eval等组合由 train.py 的main按字符串包含关系分派trainer.use_orbit: true表示使用 Orbit 自定义训练循环对应RankingTrainer置为false则走 Kerascompile/fit路径train.py 中两条路径均有完整实现fit 路径使用CheckpointManagerSimpleCheckpoint回调示例中train_steps: 256054、validation_steps: 5440、validation_interval: 85352等数值与 configs/config.py 中 Criteo TB 的全量统计一致NUM_TRAIN_EXAMPLES 4195197692、NUM_EVAL_EXAMPLES 89137318批大小 16384 时steps_per_epoch ≈ 256054评估步数89137318 // 16384 5440仓库还提供了可直接加载的 YAML 模板dlrm_criteo_tpu.yamlinteraction: dot与 dcn_v2_criteo_tpu.yamlinteraction: cross二者模型结构均为bottom_mlp: [512,256,64]、embedding_dim: 64、top_mlp: [1024,1024,512,256,1]仅交互层不同。4.2 GPU 训练GPU 训练与 TPU 流程基本相同只需更换分布策略并指定 GPU 数量以 4 卡为例。注意 GPU 上嵌入维度相应调小export EMBEDDING_DIM8 python3 official/recommendation/ranking/train.py --modetrain_and_eval \ --model_dir${BUCKET_NAME}/model_dirs/${EXPERIMENT_NAME} --params_override runtime: distribution_strategy: mirrored num_gpus: 4 ... 分布策略的解析入口是 train.py 中distribute_utils.get_distribution_strategy(distribution_strategy..., all_reduce_alg..., num_gpus..., tpu_address...)tpu策略读取runtime.tpu对应上面导出的TPU_NAME。5. 配置体系深潜所有配置以 dataclass 定义于 configs/config.py通过--params_override的 YAML 局部覆盖并有对应的单测 config_test.py。关键参数与默认值如下均为源码默认值5.1 ModelConfig模型结构参数默认值说明num_dense_features13数值特征个数须与数据列序对齐vocab_sizes[]每个稀疏特征词表大小顺序与输入数据一致embedding_dim8整数时所有表同维度列表时长度须等于vocab_sizessize_threshold50000表规模阈值小于该值用 Keras Embedding大于该值用 TPU Embedding-1表示全部用 Keras0表示全部用 TPUbottom_mlp/top_mlp[]底部/顶部 MLP 的隐藏层宽度interactiondot特征交互方式dot/cross/multi_layer_dcnconcat_denseTrue是否将交互模块输出与稠密输出再次拼接送入 top MLPdcn_num_layers3堆叠 DCN 层数dcn_low_rank_dim512堆叠 DCN 层的低秩投影维度dcn_kernel_initializer/dcn_bias_initializertruncated_normal / zerosDCN 层初始化器dcn_use_biasTrueDCN 层是否使用 biasuse_multi_hotFalse启用多热multi-hot数据加载用于 DLRM V2multi_hot_sizes[]各稀疏特征的多热大小use_partial_tpu_embeddingTrue是否使用 Partial TPU Embedding按size_threshold混合两种实现max_ids_per_chip_per_sample/max_ids_per_table/max_unique_ids_per_tableNoneSparseCore 稀疏查找的 ID 上界约束allow_id_droppingFalse是否允许丢弃超出上界的 IDinitialize_tables_on_hostFalse表超出 HBM 容量时先在 Host 内存初始化再分片拷入 HBM5.2 OptimizationConfig双优化器排序模型对嵌入权重与稠密权重使用两套独立优化器embedding_optimizer: SGD嵌入表参数dense_optimizer: AdamMLP 参数若dense_optimizer也是 SGD则使用独立的dense_sgd_config学习率曲线其warmup_steps默认为 0。5.3 LearningRateConfig 与 WarmUpAndPolyDecay嵌入优化器默认学习率曲线参数为learning_rate: 1.25、warmup_steps: 8000、decay_steps: 30000、decay_start_steps: 70000、decay_exp: 2。实际调度由 common.py 中的WarmUpAndPolyDecayKerasLearningRateSchedule实现分三段[0, warmup_steps]线性升温、[warmup_steps, decay_start_steps]恒定、[decay_start_steps, decay_start_stepsdecay_steps]多项式衰减学习率按batch_size / 2048线性缩放并有tf.maximum(0.0001, ...)与tf.maximum(0.01, ...)两处下限保护。5.4 DataConfig 与 TrainerConfigDataConfig控制输入管道input_path支持 glob、global_batch_size、dtype: float32、shuffle_buffer_size: 10000、cycle_length: 10、sharding: True、num_shards_per_host: 8、use_cached_data: False。TrainerConfig控制训练节奏train_steps、validation_steps-1 表示评估整个数据集、validation_interval每 N 步评估一次、use_orbit: False默认走 compile/fit、enable_metrics_in_training: True、time_history.log_steps以及性能开关pipeline_sparse_and_dense_execution嵌入查找与稠密计算流水化执行。回调由CallbacksConfig控制checkpoint/export、BackupAndRestore、TensorBoard、TimeHistory 四项。仓库通过exp_factory注册了三个预设实验工厂dlrm_criteodot 交互--experiment的默认值见 common.py、dcn_criteocross 交互、dlrm_dcn_v2_criteomulti_layer_dcnuse_multi_hot: Trueuse_partial_tpu_embedding: False 26 个特征的multi_hot_sizes。6. 源码级实现解析6.1 模型构建RankingTask.build_modeltask.py 的RankingTask.build_model展示了从配置到 Keras 模型的完整装配链_get_tpu_embedding_feature_config依据vocab_sizes与embedding_dim为每张表生成tf.tpu.experimental.embedding.TableConfigcombiner 为mean初始化器为标准差1/sqrt(dim)的截断正态并按表名embedding_table_%02d编号当设置了max_ids_per_*上界时还生成SparseCoreEmbeddingConfig按use_partial_tpu_embedding选择tfrs.experimental.layers.embedding.PartialTPUEmbeddingKeras/TPU 混合size_threshold分流或纯tfrs.layers.embedding.tpu_embedding_layer.TPUEmbedding按interaction分支构造DotInteraction(skip_gatherTrue)、Concatenate Cross或Concatenate MultiLayerDCN低秩投影维度由dcn_low_rank_dim指定组装tfrs.experimental.models.Rankingbottom_stack为final_activationrelu的 MLPtop_stack为final_activationsigmoid的 MLPconcat_dense控制稠密特征是否二次拼接用tfrs.experimental.optimizers.CompositeOptimizer将嵌入优化器绑定embedding_trainable_variables、稠密优化器绑定dense_trainable_variables随后model.compile。train_step/validation_step直接委托给模型自身的train_step/test_step并断言metrics model.metrics——即指标必须由 RankingModel 内部管理。6.2 数据管道分片与并行data_pipeline.py 的CriteoTsvReader是输入管道核心其__call__接收tf.distribute.InputContext按ctx.get_per_replica_batch_size将全局批大小折算为副本批大小_parse_fn使用tf.io.decode_csvfield_delim\tna_value-1解析 TSV 行输出{dense_features: [B, num_dense], sparse_features: {str(i): [B]}}与标签与 3.3 节描述的列序严格对应跨主机按ctx.input_pipeline_id对文件列表做shard切分主机内再按num_shards_per_host切分并用interleave(cycle_lengthparams.cycle_length)并行消费全程prefetch(AUTOTUNE)is_training为真时对文件流repeat()训练数据可无限复用use_multi_hot场景则由 data_pipeline_multi_hot.py 的CriteoTsvReaderMultiHotTSV 多热解析与合成数据生成或CriteoTFRecordReaderuse_tf_record_reader: true时读取预物化的 TFRecord接管分派逻辑见RankingTask.build_inputs。6.3 训练循环与日志common.py 定义了命令行旗标默认experimentdlrm_criteo、modetrain_and_eval另有seed与profile_steps。RankingTrainer.train_loop_end在每轮训练结束时除汇总 loss 与指标外会为每个优化器单独记录学习率曲线键形如SGD_0_learning_rate/Adam_1_learning_rate这正是双优化器设计在监控层面的体现。7. 小结能力关键路径训练入口与分布策略分派official/recommendation/ranking/train.py模型装配、双优化器、交互层分支official/recommendation/ranking/task.py全部配置默认值与三个预设实验official/recommendation/ranking/configs/config.py学习率旗标与 WarmUpAndPolyDecayofficial/recommendation/ranking/common.pyTSV 数据管道与合成数据official/recommendation/ranking/data/data_pipeline.py多热数据管道official/recommendation/ranking/data/data_pipeline_multi_hot.pyBeam 预处理分片、词表、编码official/recommendation/ranking/preprocessing/README.mdTPU 实验 YAML 模板dlrm_criteo_tpu.yaml、dcn_v2_criteo_tpu.yaml模块说明本文主体文档official/recommendation/ranking/README.md适用前提与限制文中 TPU 嵌入容量数据90 GB / 5.6 TB / 22.4 TB与示例步数、词表大小均针对 README 所述 Criteo Terabyte 数据与 TPU v3 系列换数据集或换 TPU 代次时需重算vocab_sizes示例值仅对 Criteo Terabyte 预处理结果有效。快速验证链路时建议先以use_synthetic_data: true跑通train_and_eval再切换到真实 TSV 数据并按预处理产出的词表回填vocab_sizes。【免费下载链接】modelsModels and examples built with TensorFlow项目地址: https://gitcode.com/GitHub_Trending/mode/models创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考