TensorFlow Models 多任务训练库 Multi-task Library 解析：配置、训练器与任务采样机制

TensorFlow Models 多任务训练库 Multi-task Library 解析配置、训练器与任务采样机制【免费下载链接】modelsModels and examples built with TensorFlow项目地址: https://gitcode.com/GitHub_Trending/mode/models本文以 multi_task.md 为骨架深入讲解 TensorFlow Models 仓库中的多任务训练库official/modeling/multitask/。该库为多任务multi-task训练与评估提供轻量级接口和通用组件不对任务类型和模型结构做任何假设只负责把多个单任务有效地“拼装”在一起。读完本文你将掌握联合joint与交替interleaving两种训练方式的源码级实现差异、MultiTaskConfig/TaskRoutine等配置对象的字段与默认值、任务采样器uniform / proportional / annealing的工作原理以及如何基于已有单任务代码如BertClassifier、BertSpanLabeler组合出共享骨干的多任务模型。一、设计定位不假设任务类型只做“任务组装脚手架”官方文档 开宗明义多任务库提供轻量级接口lite-weight interfaces与通用组件支持多任务训练和评估它不对任务类型与具体模型结构细节做假设而是设计成一种“脚手架”scaffold用于高效地组合单个任务。常见的训练调度策略实现在默认模块中并为定制化用例保留了扩展空间。文档明确该库支持两类训练模式joint联合训练各个任务各自执行前向传播forward pass得到一个联合损失joint loss随后只做一次反向传播alternative交替训练各个任务执行独立的前向与反向传播任务混合比例通过对训练步骤采样不同任务来控制。这两种模式分别对应源码中的MultiTaskBaseTrainerjoint与MultiTaskInterleavingTrainerinterleaving在 train_lib.py 中注册为训练器工厂TRAINERS { interleaving: interleaving_trainer.MultiTaskInterleavingTrainer, joint: base_trainer.MultiTaskBaseTrainer }从源码结构看trainer_type配置项的取值与上述键一一对应这也是后文配置系统中最关键的调度开关。二、组件总览接口层与通用组件层官方文档 将库分为“接口”和“通用组件”两层以下表格完整继承文档的组件划分并补充各文件在仓库中的路径与职责组件文件路径职责MultiTask接口multitask.py持有多数Task实例并保存多任务调度信息如任务权重MultiTaskBaseModel接口base_model.py暴露每个单任务的前向计算入口每个任务表示为一个tf.keras.Model实例任务间参数共享留给具体实现基础训练器base_trainer.py优化含异构数据集的多任务模型的抽象默认执行联合反向步骤可通过对不同任务损失设置权重实现平衡交替训练器interleaving_trainer.py继承基础训练器共享其损失/指标聚合与上报逻辑在任务间交替执行将单任务训练步混合在异构数据集上支持按概率分布采样任务评估器evaluator.py对每个单任务做评估的组合循环遍历指定任务用对应数据集执行评估训练驱动库train_lib.py把模型、任务、训练器组装起来并触发训练/评估执行配置configs.py提供系统顶层视图配置对象尽可能模仿或复用对应单任务组件以保持一致性三、MultiTask任务容器与联合训练步骤3.1 构造与配置入口MultiTask类是多个Task实例的“持有者”stakeholder同时保存多任务调度信息如任务权重。其构造函数接收tasksTask的列表或扁平字典。传入列表时会按task.name建字典并拒绝重名任务抛出ValueError(Duplicated tasks found, ...)task_weights任务名到权重的字典。文档强调其双重用途——既可直接作用于联合反向步骤中的损失求和也可用于交替反向步骤中采样任务未指定的任务默认权重为1.0task_eval_steps任务名到评估步数的字典name实例名。更常见的用法是从配置构造MultiTask.from_config接收MultiTaskConfig遍历其中的task_routines用task_factory.get_task逐个创建Task实例并收集每个任务的eval_steps与task_weight。这里也印证了文档中“TaskRoutine复用Task配置”的说法。另外MultiTask提供了create_optimizer类方法直接委托给Task.create_optimizer支持传入优化器配置、运行时配置以及差分隐私differential privacy配置。3.2joint_train_step一次前向多任务、一次反向联合训练步骤的源码joint_train_step完整体现了 joint 模式的语义losses {} with tf.GradientTape() as tape: total_loss 0.0 for name, model in multi_task_model.sub_tasks.items(): inputs task_inputs[name] # 支持 (features, labels) 元组或 dict 两种输入形式 outputs model(features, trainingTrue) task_loss self.tasks[name].build_losses(labels, outputs) task_weight self.task_weight(name) total_loss task_weight * task_loss losses[name] task_loss self.tasks[name].process_metrics(task_metrics[name], labels, outputs, **kwargs) # 由于默认梯度 allreduce 在优化器内部做 sum这里先缩放损失 scaled_loss total_loss / tf.distribute.get_strategy().num_replicas_in_sync tvars multi_task_model.trainable_variables grads tape.gradient(scaled_loss, tvars) optimizer.apply_gradients(list(zip(grads, tvars))) losses[total_loss] total_loss return losses几个值得注意的实现细节加权求和每个任务的损失通过task_weight(name)加权后累加到total_loss这正是文档所说“通过设置不同任务权重来平衡任务”在 joint 模式下的落地方式复用单任务的损失与指标build_losses与process_metrics均调用自各Task实例多任务库没有另起炉灶定义损失函数——这与文档“完全复用输入函数、损失函数、指标及其聚合与归约逻辑”的表述一致分布式缩放scaled_loss除以num_replicas_in_sync源码注释说明原因是默认梯度 allreduce 在优化器内部执行求和返回值返回包含各任务损失及total_loss的字典供训练器的指标记录使用。四、MultiTaskBaseModel任务间参数共享留给实现文档 指出base_model.py提供每个单任务前向计算的访问入口每个任务表示为一个tf.keras.Model实例任务间参数共享交给具体实现。源码 的结构印证了这一点MultiTaskBaseModel.__init__调用抽象方法_instantiate_sub_tasks()返回“任务名 -tf.keras.Model”的映射未实现时抛出NotImplementedErrorsub_tasks属性暴露该映射供训练器/评估器取出某个任务的子模型注意joint_train_step与评估器中都是按multi_task_model.sub_tasks[name]索引前向计算的initialize()是可选钩子用于加载预训练 checkpointbuild()逐个调用未构建的子模型task_model(task_model.inputs)以确保变量全部创建——这一步对新版 Keras 优化器是必需的因为 交替训练器 在检测到tf_keras.optimizers.experimental.Optimizer时会显式调用multi_task_model.build()并optimizer.build(trainable_variables)让优化器一次性创建所有 slot 变量。单任务的可组合性以 BERT 分类 序列标注为例文档 的“Notes on single task composability”一节强调库的设计目标是通过组合已有的单任务实现来搭建多任务模型。给出的范例是BertClassifier与BertSpanLabeler后者的骨干网络backbone直接从 classifier 对象初始化因此“分类 序列标注、共享骨干编码器”的多任务模型可以从现有代码轻松构造——共享部分通过构造函数传入即可复用单任务的tf.keras.Model实现。文档还特别警示了一个易错点在多任务训练场景下单任务的build_model()不会被使用。原因是部分共享结构partially shared structure无法仅由单个任务来指定——模型必须由MultiTaskBaseModel统一实例化。这与 train_lib.py 中的分支逻辑一致当任务是MultiTask时走build_multitask_model()普通任务才走build_model()。五、两种训练器的源码级对比5.1MultiTaskBaseTrainerjointMultiTaskBaseTrainer继承 Orbit 的StandardTrainer构造时通过tf.distribute.get_strategy()获取当前分发策略若优化器是ExponentialMovingAverage且尚未创建影子权重则调用shadow_copy(multi_task_model)创建权重移动平均的影子副本将model、optimizer、global_step以及模型的checkpoint_items若存在打包进tf.train.Checkpoint若未显式传入train_datasets则按任务名自动构建每个任务用各自的task.build_inputs与task_config.train_data经orbit.utils.make_distributed_dataset包装成分布式数据集——这正是文档所说“涉及异构数据集heterogeneous datasets”的抽象所在。其train_step在策略范围内对“任务名 - 迭代器”的映射做next一次性取出所有任务的 batch然后调用multi_task.joint_train_step执行第三节描述的加权联合反向步骤并把返回的各任务损失与total_loss记入training_losses指标。此外training_losses会为total_loss和每个任务各建一个Mean指标training_metrics则调用每个Task的build_metrics(trainingTrue)。train_loop_end在每轮循环结束时按任务输出指标并记录学习率源码注释指出学习率调度由 Keras 优化器内部管理遵循的是反向传播次数iterations而不是多任务训练器的逻辑 global step。5.2MultiTaskInterleavingTraineralternative / interleavingMultiTaskInterleavingTrainer继承基础训练器因而共享其损失、指标聚合与上报等“家务”逻辑与文档描述一致差异在于训练步骤的调度方式按任务预构建训练步构造时为每个任务构建step_fn——若模型是MultiTaskBaseModel则取sub_tasks[task_name]否则直接复用整个模型允许单任务模型被拿来在多任务框架下交替训练每个任务的步调用的是该任务自己的train_step()从而把优化细节隐藏在单个任务内部训练器只关注调度与任务平衡这正是文档第三点的含义维护 per-task 步计数器task_step_counter(name)为每个任务建立独立的 global step 变量便于后续统计各任务实际训练了多少步源码中标注了将用于 TensorBoard 日志的 TODO核心采样逻辑在train_stepdef train_step(self, iterator_map): # 按多项式分布采样一个任务进行训练 rn tf.random.stateless_uniform(shape[], seed(0, self.global_step)) cumulative_sample_distribution self._task_sampler.task_cumulative_distribution( self.global_step) # 前置 [0.0] 便于索引 cumulative_sample_distribution tf.concat( [tf.constant([0.0], dtypetf.float32), cumulative_sample_distribution], axis0) for idx, (name, _) in enumerate(self.multi_task.tasks.items()): begin cumulative_sample_distribution[idx] end cumulative_sample_distribution[idx 1] if rn begin and rn end: self._strategy.run( self._task_train_step(name), args(next(iterator_map[name]),))可以看到每个训练步先用以global_step为种子的无状态随机数stateless uniform保证分布式各副本采样一致抽取一个[0, 1)区间随机数再与任务采样器给出的累积分布做区间匹配命中哪个任务就只消费该任务数据集的下一个 batch 并执行其单任务train_step。命中的任务同时累加global_step与自己的task_step_counter。另一个细节interleaving 模式下total_loss没有语义源码注释说明其值恒为零因此train_loop_end会把total_loss从日志结果中剔除。这一点也被 测试用例 明确断言训练结束后各任务结果键中应包含training_loss与任务指标如bar_acc、foo_acc且total_loss不在结果中。六、任务采样器三种分布控制任务混合比例文档提到“将支持按概率分布采样任务以促进任务平衡”这部分在 task_sampler.py 中已完整实现。抽象基类TaskSampler只有一个抽象方法task_cumulative_distribution(global_step)返回形状为(num_tasks, 1)的累积分布张量供第五节的区间匹配采样使用。内置三种采样器采样器行为配置对象UniformTaskSampler所有任务均匀采样即每个任务概率1/Nuniform分支ProportionalTaskSampler按任务权重的alpha次幂归一化后采样p_i w_i^alpha / Σ w_j^alphaProportionalSampleConfigalpha: float 1.0AnnealingTaskSampler采样分布随训练进度退火alpha从 1.0 线性衰减到 0.2实现从“前期偏向大数据量任务”到“后期近似均匀”的过渡AnnealingSampleConfigsteps_per_epoch: int 5total_steps: int 20其中ProportionalTaskSampler的alpha幂函数意味着alpha越大权重大的任务被采样的相对占比越高alpha1.0时即严格按任务权重比例采样。AnnealingTaskSampler的alpha随 epoch 推进按1.0 - 0.8 * (cur_epoch - 1) / (total_epochs - 1 1e-10)动态变化源码第 102 行其设计参考了源码注释中引用的多任务学习采样策略论文Stickland et al., 2019。采样器由工厂函数get_task_sampler根据TaskSamplingConfig.typeuniform/proportional/annealing创建类型不识别时抛出RuntimeError。它在训练器创建链路中被 get_trainer 使用仅当params.trainer.trainer_type interleaving时才构造采样器并注入训练器且采样权重直接取自task.task_weights。七、配置系统从TaskRoutine到实验配置文档 强调configs.py提供“顶层视图”配置对象尽可能模仿或复用对应单任务组件。源码 中的配置层级如下配置对象字段含默认值说明TaskRoutinetask_name: str 、task_config: TaskConfig \| None None、eval_steps: Optional[int]、task_weight: Optional[float] 1.0单个任务的例行配置复用单任务的TaskConfigMultiTaskConfiginit_checkpoint: str 、model: Config \| None None、task_routines: Tuple[TaskRoutine, ...] ()、differential_privacy_config已弃用多任务顶层配置角色类似单任务的TaskConfigProportionalSampleConfigalpha: float 1.0proportional 采样的指数AnnealingSampleConfigsteps_per_epoch: int 5、total_steps: int 20退火采样的 epoch 步数与总步数TaskSamplingConfigtype: str 以及uniform/proportional/annealing三个 OneOf 分支OneOf 配置type决定生效分支MultiTaskTrainerConfigtrainer_type: str interleaving、task_sampler: TaskSamplingConfig默认typeproportional继承TrainerConfig选择 joint 或 interleaving 训练器MultiTaskExperimentConfigtask: MultiTaskConfig、trainer: MultiTaskTrainerConfig、runtime: RuntimeConfig多任务训练/评估的实验配置MultiEvalExperimentConfigeval_tasks: Tuple[TaskRoutine, ...] ()单任务训练 多任务评估的实验配置下面是一个可直接参考的配置构造示例取自 interleaving_trainer_test.py 的真实测试config configs.MultiTaskConfig( task_routines( configs.TaskRoutine( task_namefoo, task_configtest_utils.FooConfig(), task_weight3.0), configs.TaskRoutine( task_namebar, task_configtest_utils.BarConfig(), task_weight1.0))) test_multitask multitask.MultiTask.from_config(config)示例中foo任务权重为 3.0、bar为 1.0在 joint 训练中foo的损失贡献三倍在 interleaving 训练中默认typeproportional、alpha1.0foo被采样的概率约为3/4。这与 multitask.py 中对task_weights的双重用途说明完全对应。八、评估器与训练驱动多任务如何被“跑起来”8.1MultiTaskEvaluatorMultiTaskEvaluator实现 Orbit 的AbstractEvaluator接口正如文档所说“对每个单任务评估做组合循环遍历指定任务并用对应数据集评估”。其实现要点每个任务用自己的task.build_inputs与task_config.validation_data构建独立评估数据集每个任务封装为独立tf.function循环调用task.validation_step前向并更新该任务的Mean损失指标模型同样支持MultiTaskBaseModel取sub_tasks[task_name]或普通单任务模型evaluate先重置所有任务的损失与指标再逐任务运行eval_steps步每任务步数可经TaskRoutine.eval_steps单独指定缺省时回退到传入的num_steps用task.aggregate_logs/task.reduce_aggregated_logs完成自定义聚合例如 BLEU 类需要跨步累积的指标最终返回“任务名 - 指标字典”的结构化结果评估结束时若配置了checkpoint_exporter会调用maybe_export_checkpoint导出最优 checkpoint。8.2train_lib.run_experiment训练驱动入口run_experiment是文档中“把模型、任务、训练器组装起来并触发训练/评估执行”的落点。关键流程若处于训练模式且未传入trainer则用get_trainer按params.trainer.trainer_type创建对应训练器joint 或 interleavinginterleaving 时附带任务采样器若处于评估模式创建MultiTaskEvaluator其eval_steps直接取自task.task_eval_steps用tf.train.CheckpointManager管理 checkpointmax_to_keep、checkpoint_interval来自params.trainer并以model.initialize作为init_fn——即MultiTaskBaseModel的预训练 checkpoint 加载钩子在此被接上构造orbit.Controller含训练/验证 summary 目录model_dir/train与model_dir/validation按mode分派mode取值执行行为traincontroller.train(stepsparams.trainer.train_steps)train_and_eval按validation_interval交替训练train_steps步与验证validation_steps步evalcontroller.evaluate(stepsparams.trainer.validation_steps)continuous_eval持续评估直至global_step train_stepstimeout_fn判定或超时run_post_evalTrue时额外执行一次完整评估并返回指标字典否则只返回模型。8.3 单任务训练 多任务评估除全量多任务训练外run_experiment_with_multitask_eval 支持另一种组合训练侧是单个Task走核心库的core_lib.Trainer评估侧是多个Task列表配合MultiEvalExperimentConfig.eval_tasks使用。从源码结构看这一路径适合“先用某个主任务训练骨干、再在多个下游任务上统一评估”的场景且支持通过train_actions_factory/eval_actions_factory注入自定义 Orbit 动作如新最优指标动作、抢占时保存 checkpoint 等进一步扩展了文档所说的“为定制化用例保留扩展空间”。九、要点回顾与使用建议模式选择需要任务梯度在同一反向步骤中相互影响的场景选trainer_typejoint损失按task_weight加权见 joint_train_step希望各任务独立反向、避免批次间耦合、并可灵活控制混合比例的场景选trainer_typeinterleaving默认值通过task_sampler.type与TaskRoutine.task_weight控制采样分布参数共享共享结构必须经由MultiTaskBaseModel的子模型构造完成单任务的build_model()在多任务下不生效参考BertClassifier/BertSpanLabeler的骨干复用方式日志语义差异joint 模式日志含total_lossinterleaving 模式有意移除total_loss见 train_loop_end 及测试断言阅读 TensorBoard 曲线时应注意这一点分布式前提训练需在tf.distribute.Strategy作用域内构建训练器与数据集测试用例 覆盖了默认、单 GPU、云 TPU 三种策略组合损失在joint_train_step中已按同步副本数缩放。以上组件、字段与行为均可在 official/modeling/multitask/ 目录下的源码与配套测试*_test.py中逐一验证结合 official/nlp/docs/multi_task.md 的设计说明构成了该仓库多任务训练能力的完整参考。【免费下载链接】modelsModels and examples built with TensorFlow项目地址: https://gitcode.com/GitHub_Trending/mode/models创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考