TensorFlow Official Models 优化配置实战OptimizerFactory、学习率衰减与 Warmup 调度全解析【免费下载链接】modelsModels and examples built with TensorFlow项目地址: https://gitcode.com/GitHub_Trending/mode/models本文基于 TensorFlow Official ModelsTFM仓库中的优化器与学习率调度文档 optimization.md系统讲解official/modeling/optimization包的配置驱动式训练优化机制如何通过一份声明式配置构建优化器、学习率衰减与 warmup 调度warmup 与衰减在数值上如何衔接以及如何在 Task 层定制优化器。读完本文你可以直接复用仓库现成的优化配置体系完成模型训练调参并能读懂 optimizer_factory.py 的底层调用链与配置数据类的每个默认值。优化包的定位与三步构建流程TFM 的 optimization 包 统一管理训练中的两大要素优化器Optimizer与学习率调度Learning Rate Scheduler。其核心是一个工厂类OptimizerFactory它接收一份优化配置作为输入提供成员函数分别构建学习率调度和优化器实例。按照文档定义的标准流程创建一个SGD 优化器 阶梯式学习率衰减 线性 warmup只需三步定义优化配置包含优化器、学习率调度以及可选的 warmup用该配置初始化OptimizerFactory实例调用类成员函数构建学习率再构建优化器。文档给出的完整示例此处将原示例的third_party.tensorflow_models.official.modeling导入前缀改写为本仓库内的等价路径official.modeling两者指向同一模块from official.modeling import optimization params {optimizer: { type: sgd, sgd: {momentum: 0.9}}, learning_rate: {type: stepwise, stepwise: { boundaries: [10000, 20000], values: [0.1, 0.01, 0.001]}}, warmup: {type: linear, linear: {warmup_steps: 500, warmup_learning_rate: 0.01}}} # Defines optimization config from a dictionary. opt_config optimization.OptimizationConfig(params) # Initializes an optimization factory from optimization config. opt_factory optimization.OptimizerFactory(opt_config) # Builds the desired learning rate scheduling instance. lr opt_factory.build_learning_rate() # Builds the optimizer instance with the desired learning rate schedule. optimizer opt_factory.build_optimizer(lr)这段示例与工厂类的 docstring 完全一致可参见 optimizer_factory.py。配置侧的约束是optimizer与learning_rate字段必须定义warmup为可选字段每个组件通过type字段指定具体类型oneof 语义。从源码结构看配置到实例的映射由两层代码完成OptimizationConfig 是顶层 dataclass聚合了optimizer、ema、learning_rate、warmup四个 oneof 子配置OptimizerFactory.__init__在初始化时即解析出self._optimizer_type、self._lr_type、self._warmup_type并在optimizer或learning_rate类型缺失时直接抛出ValueError见 optimizer_factory.py#L140-L163这保证了配置错误在构建阶段就暴露而不是在训练循环中才报错。oneof 机制由 oneof.py 实现OneOfConfig要求type字段恰好指向一个已定义的子字段从而保证type 与实际配置段始终一一对应。支持的优化器及其配置字段文档列出的基础支持列表对应早期版本的OPTIMIZERS_CLS为OPTIMIZERS_CLS { sgd: tf.keras.optimizers.SGD, adam: tf.keras.optimizers.Adam, adamw: nlp_optimization.AdamWeightDecay, lamb: tfa_optimizers.LAMB, rmsprop: tf.keras.optimizers.RMSprop }当前仓库的 optimizer_factory.py#L31-L57 在此基础上进一步扩展LEGACY_OPTIMIZERS_CLS提供sgd、adam、rmsprop、adagrad等 legacy 实现NEW_OPTIMIZERS_CLS提供同名 experimental新版 Keras实现SHARED_OPTIMIZERS则统一挂载sgd_experimental、adam_experimental、adamw、adamw_experimental、lamb、lars、slide、adafactor、adafactor_keras。也就是说现在可选的type值远多于文档列举的五种build_optimizer通过use_legacy_optimizer参数默认True在两套字典间切换。各类型优化器的可用配置字段含默认值定义在 optimizer_config.py几个常用的 dataclass 及其字段如下配置类关键字段默认值说明SGDConfigmomentum: 0.0、nesterov: False、decay: 0.0与 Keras SGD 参数一致AdamConfigbeta_1: 0.9、beta_2: 0.999、epsilon: 1e-07、amsgrad: False支持 AMSGrad 变体RMSPropConfigrho: 0.9、momentum: 0.0、epsilon: 1e-7、centered: Falserho即折扣因子AdamWeightDecayConfigweight_decay_rate: 0.0、include/exclude_from_weight_decay、gradient_clip_norm: 1.0可指定权重衰减的纳入/排除名单LARSConfigmomentum: 0.9、eeta: 0.001、weight_decay_rate: 0.0、classic_momentum: True层自适应学习率大 batch 训练常用EMAConfigaverage_decay: 0.99、start_step: 0、dynamic_decay: True配置ema字段后自动套一层指数滑动平均梯度裁剪所有优化器都支持三种梯度裁剪方式按值裁剪clip by value、按范数裁剪clip by norm、按全局范数裁剪clip by global norm。对应字段定义在 BaseOptimizerConfig#L34-L36clipnorm、clipvalue、global_clipnorm均为None时表示不启用。在build_optimizer的实现中optimizer_factory.py#L222-L236三个裁剪字段为None时会被先从参数字典中删除非None的则作为关键字参数透传给底层 Keras 优化器构造函数。另外需要注意一个版本差异使用新版 Keras 优化器use_legacy_optimizerFalse时配置decay字段会直接报错因为新版优化器已弃用decay参数需要把衰减逻辑并入lr调度或改用 legacy 优化器见 optimizer_factory.py#L241-L245。示例RMSProp 全局范数裁剪文档给出的例子是指定一个折扣因子rho为 0.9、全局范数裁剪阈值为 10.0 的 RMSProp 优化器params {optimizer: { type: rmsprop, rmsprop: {rho: 0.9, global_clipnorm: 10.0}}}由于rho的默认值恰好就是 0.9这里显式写出的意义在于自描述真正改变行为的是global_clipnorm: 10.0它会让所有权重梯度的全局 L2 范数不超过 10.0。如何添加一个新的优化器文档给出的扩展流程分三步创建自定义优化器即 继承tf_keras.optimizers.Optimizer的子类在 optimization/configs/optimizer_config.py 中为新优化器添加所需的配置 dataclass 字段并在 OptimizationConfig 的 OptimizerConfig 里挂上对应的 oneof 字段OptimizationConfig目前暴露sgd、sgd_experimental、adam、adam_experimental、adamw、adamw_experimental、lamb、rmsprop、lars、adagrad、slide、adafactor、adafactor_keras等全部子配置把优化器类加入工厂的可用类列表即 optimizer_factory.py 中的优化器字典。当前仓库额外提供了运行时注册入口 register_optimizer_clsregister_optimizer_cls(key, optimizer_config_cls, use_legacy_optimizer)可以把自定义优化器类动态挂到LEGACY_OPTIMIZERS_CLS或NEW_OPTIMIZERS_CLS重复注册会抛ValueError。其 docstring 也明确提示仅注册类还不够用户仍需为它提供 oneof 配置 dataclass 才能被OptimizerFactory识别——这与文档三步流程中的第 2 步相互印证。学习率与 Warmup 调度学习率通过learning_rate字段必填配置warmup 通过warmup字段可选配置。文档列出的调度类型清单LR_CLS { stepwise: tf.keras.optimizers.schedules.PiecewiseConstantDecay, polynomial: tf.keras.optimizers.schedules.PolynomialDecay, exponential: tf.keras.optimizers.schedules.ExponentialDecay, cosine: tf.keras.experimental.CosineDecay, power: lr_schedule.DirectPowerDecay, } WARMUP_CLS { linear: lr_schedule.LinearWarmup, polynomial: lr_schedule.PolynomialWarmUp }此外还允许constant常数学习率类型从源码看optimizer_factory.py#L176-L179当type constant时build_learning_rate直接返回lr_config.learning_rate这个浮点值不构造 schedule 对象。当前仓库的 LR_CLS#L59-L69 已扩充为stepwise、polynomial、exponential、cosine、cosine_restarts、power、power_linear、power_with_offset、step_cosine_with_offset共九种其中前五种衰减类型均带WithOffset后缀是由 _make_offset_wrapper 动态生成的子类——行为等价于new_class_object(step) base_lr_class_object(step - offset)即给所有调度统一加上了一个offset步数偏移能力各配置 dataclass 中的offset字段默认 0 即为由此而来。各调度的配置字段与默认值定义在 learning_rate_config.py常用项摘录配置类关键字段默认值说明ConstantLrConfiglearning_rate: 0.1常数学习率StepwiseLrConfigboundaries、values均为必填、offset: 0values必须比boundaries多一个元素[0, b0] - v0、[b0, b1] - v1……ExponentialLrConfiginitial_learning_rate、decay_steps、decay_rate必填、staircase: False、offset: 0指数衰减可开启阶梯模式PolynomialLrConfigend_learning_rate: 0.0001、power: 1.0、cycle: False多项式衰减cycleTrue时可循环CosineLrConfiginitial_learning_rate、decay_steps必填、alpha: 0.0余弦衰减alpha是末端学习率占初始值比例CosineRestartsLrConfigfirst_decay_steps必填、t_mul: 2.0、m_mul: 1.0、alpha: 0.0余弦重启周期逐段拉长/重置DirectPowerLrConfigpower: -0.5遵循lr * step^power默认即平方根衰减PowerAndLinearDecayLrConfigtotal_decay_steps必填、power: -0.5、linear_decay_fraction: 0.1先 power 衰减末尾一段线性压到 0LinearWarmupConfigwarmup_learning_rate: 0、warmup_steps必填线性 warmup 起点学习率PolynomialWarmupConfigpower: 1、warmup_steps必填多项式 warmup学习率工作机制Warmup 与衰减如何拼接文档对学习率如何工作给出了三条明确的规则这是理解整套调度语义的关键学习率函数以step为输入返回该步的学习率值随训练推进学习率通常逐渐衰减warmup 调度常用于稳定训练过程它从一个较低的学习率起步逐步增大直至达到常规衰减调度的初始值拼接方式如下设 warmup 步数为warmup_steps总训练步数为train_steps步骤[0, warmup_steps)learning_rate warmup(step)步骤[warmup_steps, train_steps)learning_rate lr(step)warmup 的终点学习率不是用户指定的而是由衰减调度反推得到的即learning_rate(warmup_steps) warmup(warmup_steps)。并且 warmup不会把常规衰减整体推迟 warmup_steps 步而是替换掉这段区间的衰减。LinearWarmup 的实现与上述语义一一对应。构造时若主调度是 schedule 对象则执行self._final_warmup_lr after_warmup_lr_sched(warmup_steps)L131-L135——这正是终点值由衰减调度在warmup_steps处求值反推的源码体现。每一步的取值逻辑为linear_warmup_lr ( self._init_warmup_lr global_step / self._warmup_steps * (self._final_warmup_lr - self._init_warmup_lr)) lr tf.cond(global_step self._warmup_steps, lambda: linear_warmup_lr, lambda: after_warmup_lr) # 主调度按原始 step 求值不做偏移注意after_warmup_lr用的是未经偏移的原始step而非step - warmup_steps这从代码层面印证了替换而非推迟的设计warmup 结束后衰减曲线的位置与没有 warmup 时完全一致只是[0, warmup_steps)区间被线性爬升覆盖。同理 PolynomialWarmUp 中warmup_learning_rate initial_learning_rate * (step / warmup_steps)^power且对step0做了tf.math.maximum(step, 1.0)保护以避免除零产生 Inf。关于学习率日志的可见性文档还提醒学习率数值是按 summary_intervalTrainerConfig.summary_interval即两次 summary 写入之间的步数周期性记录的。如果warmup_steps小于summary_interval你在 summary 里就看不到 warmup 阶段的取值。以 imagenet_resnet50_tpu.yaml 为例summary_interval: 312而warmup_steps: 15605 个采样点warmup 曲线完整可见反之若把 warmup 设成 100 步则大概率整个 warmup 过程都不会出现在学习率曲线上——调参观察时应留意这一点。示例Cosine 衰减 线性 Warmup文档给出的另一个例子指定 decay_steps 为 20000 的余弦学习率衰减并对前 500 步做线性 warmupparams {learning_rate: {type: cosine, cosine: {decay_steps: 20000}}, warmup: {type: linear, linear: {warmup_steps: 500}}}按前文的拼接规则[0, 500)区间内学习率从warmup_learning_rate默认 0线性爬升到cosine(500)处反推出的值[500, 20000)则严格沿余弦曲线衰减曲线位置不因 warmup 而后移。实战视角视觉实验配置中的 optimizer_config在 TFM 视觉模块的实验配置里优化配置统一挂在trainer段的optimizer_config字段下。imagenet_resnet50_tpu.yaml#L27-L47 是一份可直接对照文档示例的完整实例trainer: train_steps: 28080 validation_steps: 13 validation_interval: 312 steps_per_loop: 312 summary_interval: 312 checkpoint_interval: 312 optimizer_config: optimizer: type: sgd sgd: momentum: 0.9 learning_rate: type: stepwise stepwise: boundaries: [9360, 18720, 24960] values: [1.6, 0.16, 0.016, 0.0016] warmup: type: linear linear: warmup_steps: 1560这份配置恰好演示了文档重要考量两点的实际落地train steps 与衰减边界的联动boundaries的 9360 / 18720 / 24960 分别是train_steps28080 的 1/3、2/3 和 5/7——衰减边界是相对总步数按比例设定的而不是拍脑袋的绝对值batch size 与学习率的联动该实验global_batch_size: 4096对应初始学习率 1.61/4096 规模的线性缩放结果warmup 步数 1560 约等于 10 个 epoch 内的一个固定爬升段。若把全局 batch size 减半或翻倍学习率各档位1.6 → 0.16 → 0.016 → 0.0016与训练步数都需要同步重估。optimizer_config这一字段在训练框架中的承接者是 TrainerConfig它与train_steps、summary_interval等调度类参数同处一个配置段保证训练节奏与优化节奏在一份配置中集中维护。在 Task 中定制优化器文档指出优化器与学习率的创建发生在 Task 内部如果任务需要不同的优化器/学习率调度器可以通过覆写类方法实现。对应源码是 Task.create_optimizerclassmethodclassmethod def create_optimizer(cls, optimizer_config: OptimizationConfig, runtime_config: Optional[RuntimeConfig] None, dp_config: Optional[DifferentialPrivacyConfig] None): gradient_transformers None if dp_config is not None: noise_stddev dp_config.clipping_norm * dp_config.noise_multiplier gradient_transformers [ functools.partial(ops.clip_l2_norm, l2_norm_clipdp_config.clipping_norm), functools.partial(ops.add_noise, noise_stddevnoise_stddev) ] opt_factory optimization.OptimizerFactory(optimizer_config) optimizer opt_factory.build_optimizer( opt_factory.build_learning_rate(), gradient_transformersgradient_transformers) if runtime_config: optimizer performance.configure_optimizer( optimizer, use_float16runtime_config.mixed_precision_dtype float16, loss_scaleruntime_config.loss_scale) return optimizer从源码结构看这个默认实现在构建优化器这一核心动作之外还叠加了三项框架级能力自定义任务覆写时可以按需保留或替换差分隐私梯度变换当传入dp_config时向build_optimizer注入clip_l2_normadd_noise两个 gradient transformer这正是 build_optimizer 签名 中gradient_transformers参数的用途且此时不应再设置clipvalue/clipnorm/global_clipnorm二者互斥混合精度配置若 runtime 配置了mixed_precision_dtype与loss_scale如 ResNet50 实验中的bfloat16会用performance.configure_optimizer为优化器套上 loss scale避免 float16 下溢出/下溢EMA 包装若OptimizationConfig.ema非空build_optimizer内部会自动把优化器包进 ExponentialMovingAverage要求use_legacy_optimizerTrue见 optimizer_factory.py#L248-L254。因此视觉/检测等项目里自定义优化器的标准姿势是在任务类中覆写create_optimizer这个类方法复用OptimizerFactory处理配置到实例的映射只在自己需要的位置如自定义梯度变换、额外的后处理postprocessor插入差异逻辑而不必从零实现调度逻辑。构建结果的正确性有 optimizer_factory_test.py 等测试文件覆盖可运行仓库测试验证你的覆写是否与默认行为一致。重要考量因素文档最后强调了两个调参时必须成对考虑的因素这里结合前面的配置实例再作总结Batch size改变 batch size 通常需要同比例缩放学习率数值和训练步数步数与样本吞吐相关。修改 batch size 后请同步调整learning_rate各档位数值与train_steps只改其一会破坏原有的等效训练预算Train steps训练总步数与decay_stepscosine、boundariesstepwise、total_decay_stepspower_linear等字段高度相关。文档明确警告只改其中一个不改另一个可能导致非预期的行为。以 imagenet_resnet50_tpu.yaml 为例若把train_steps从 28080 翻倍却不动boundaries: [9360, 18720, 24960]衰减将在总步数不到一半时就全部走完后半程训练将以最低档学习率空转——这就是文档所指的非预期行为。关键文件索引内容仓库路径本文对应的官方文档optimization.md优化器工厂构建调度与优化器、类型注册optimizer_factory.py优化器配置 dataclass含裁剪字段、默认值optimizer_config.py学习率/warmup 配置 dataclass含默认值learning_rate_config.py顶层 OptimizationConfigoptimizer/ema/learning_rate/warmup 聚合optimization_config.py调度实现LinearWarmup、PolynomialWarmUp、Power 系衰减、offset 包装lr_schedule.pyoneof 配置语义oneof.pyTask.create_optimizer 默认实现DP、混合精度base_task.pyTrainerConfigsummary_interval 等训练节奏参数config_definitions.py完整实验配置示例imagenet_resnet50_tpu.yaml【免费下载链接】modelsModels and examples built with TensorFlow项目地址: https://gitcode.com/GitHub_Trending/mode/models创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
TensorFlow Official Models 优化配置实战：OptimizerFactory、学习率衰减与 Warmup 调度全解析
TensorFlow Official Models 优化配置实战OptimizerFactory、学习率衰减与 Warmup 调度全解析【免费下载链接】modelsModels and examples built with TensorFlow项目地址: https://gitcode.com/GitHub_Trending/mode/models本文基于 TensorFlow Official ModelsTFM仓库中的优化器与学习率调度文档 optimization.md系统讲解official/modeling/optimization包的配置驱动式训练优化机制如何通过一份声明式配置构建优化器、学习率衰减与 warmup 调度warmup 与衰减在数值上如何衔接以及如何在 Task 层定制优化器。读完本文你可以直接复用仓库现成的优化配置体系完成模型训练调参并能读懂 optimizer_factory.py 的底层调用链与配置数据类的每个默认值。优化包的定位与三步构建流程TFM 的 optimization 包 统一管理训练中的两大要素优化器Optimizer与学习率调度Learning Rate Scheduler。其核心是一个工厂类OptimizerFactory它接收一份优化配置作为输入提供成员函数分别构建学习率调度和优化器实例。按照文档定义的标准流程创建一个SGD 优化器 阶梯式学习率衰减 线性 warmup只需三步定义优化配置包含优化器、学习率调度以及可选的 warmup用该配置初始化OptimizerFactory实例调用类成员函数构建学习率再构建优化器。文档给出的完整示例此处将原示例的third_party.tensorflow_models.official.modeling导入前缀改写为本仓库内的等价路径official.modeling两者指向同一模块from official.modeling import optimization params {optimizer: { type: sgd, sgd: {momentum: 0.9}}, learning_rate: {type: stepwise, stepwise: { boundaries: [10000, 20000], values: [0.1, 0.01, 0.001]}}, warmup: {type: linear, linear: {warmup_steps: 500, warmup_learning_rate: 0.01}}} # Defines optimization config from a dictionary. opt_config optimization.OptimizationConfig(params) # Initializes an optimization factory from optimization config. opt_factory optimization.OptimizerFactory(opt_config) # Builds the desired learning rate scheduling instance. lr opt_factory.build_learning_rate() # Builds the optimizer instance with the desired learning rate schedule. optimizer opt_factory.build_optimizer(lr)这段示例与工厂类的 docstring 完全一致可参见 optimizer_factory.py。配置侧的约束是optimizer与learning_rate字段必须定义warmup为可选字段每个组件通过type字段指定具体类型oneof 语义。从源码结构看配置到实例的映射由两层代码完成OptimizationConfig 是顶层 dataclass聚合了optimizer、ema、learning_rate、warmup四个 oneof 子配置OptimizerFactory.__init__在初始化时即解析出self._optimizer_type、self._lr_type、self._warmup_type并在optimizer或learning_rate类型缺失时直接抛出ValueError见 optimizer_factory.py#L140-L163这保证了配置错误在构建阶段就暴露而不是在训练循环中才报错。oneof 机制由 oneof.py 实现OneOfConfig要求type字段恰好指向一个已定义的子字段从而保证type 与实际配置段始终一一对应。支持的优化器及其配置字段文档列出的基础支持列表对应早期版本的OPTIMIZERS_CLS为OPTIMIZERS_CLS { sgd: tf.keras.optimizers.SGD, adam: tf.keras.optimizers.Adam, adamw: nlp_optimization.AdamWeightDecay, lamb: tfa_optimizers.LAMB, rmsprop: tf.keras.optimizers.RMSprop }当前仓库的 optimizer_factory.py#L31-L57 在此基础上进一步扩展LEGACY_OPTIMIZERS_CLS提供sgd、adam、rmsprop、adagrad等 legacy 实现NEW_OPTIMIZERS_CLS提供同名 experimental新版 Keras实现SHARED_OPTIMIZERS则统一挂载sgd_experimental、adam_experimental、adamw、adamw_experimental、lamb、lars、slide、adafactor、adafactor_keras。也就是说现在可选的type值远多于文档列举的五种build_optimizer通过use_legacy_optimizer参数默认True在两套字典间切换。各类型优化器的可用配置字段含默认值定义在 optimizer_config.py几个常用的 dataclass 及其字段如下配置类关键字段默认值说明SGDConfigmomentum: 0.0、nesterov: False、decay: 0.0与 Keras SGD 参数一致AdamConfigbeta_1: 0.9、beta_2: 0.999、epsilon: 1e-07、amsgrad: False支持 AMSGrad 变体RMSPropConfigrho: 0.9、momentum: 0.0、epsilon: 1e-7、centered: Falserho即折扣因子AdamWeightDecayConfigweight_decay_rate: 0.0、include/exclude_from_weight_decay、gradient_clip_norm: 1.0可指定权重衰减的纳入/排除名单LARSConfigmomentum: 0.9、eeta: 0.001、weight_decay_rate: 0.0、classic_momentum: True层自适应学习率大 batch 训练常用EMAConfigaverage_decay: 0.99、start_step: 0、dynamic_decay: True配置ema字段后自动套一层指数滑动平均梯度裁剪所有优化器都支持三种梯度裁剪方式按值裁剪clip by value、按范数裁剪clip by norm、按全局范数裁剪clip by global norm。对应字段定义在 BaseOptimizerConfig#L34-L36clipnorm、clipvalue、global_clipnorm均为None时表示不启用。在build_optimizer的实现中optimizer_factory.py#L222-L236三个裁剪字段为None时会被先从参数字典中删除非None的则作为关键字参数透传给底层 Keras 优化器构造函数。另外需要注意一个版本差异使用新版 Keras 优化器use_legacy_optimizerFalse时配置decay字段会直接报错因为新版优化器已弃用decay参数需要把衰减逻辑并入lr调度或改用 legacy 优化器见 optimizer_factory.py#L241-L245。示例RMSProp 全局范数裁剪文档给出的例子是指定一个折扣因子rho为 0.9、全局范数裁剪阈值为 10.0 的 RMSProp 优化器params {optimizer: { type: rmsprop, rmsprop: {rho: 0.9, global_clipnorm: 10.0}}}由于rho的默认值恰好就是 0.9这里显式写出的意义在于自描述真正改变行为的是global_clipnorm: 10.0它会让所有权重梯度的全局 L2 范数不超过 10.0。如何添加一个新的优化器文档给出的扩展流程分三步创建自定义优化器即 继承tf_keras.optimizers.Optimizer的子类在 optimization/configs/optimizer_config.py 中为新优化器添加所需的配置 dataclass 字段并在 OptimizationConfig 的 OptimizerConfig 里挂上对应的 oneof 字段OptimizationConfig目前暴露sgd、sgd_experimental、adam、adam_experimental、adamw、adamw_experimental、lamb、rmsprop、lars、adagrad、slide、adafactor、adafactor_keras等全部子配置把优化器类加入工厂的可用类列表即 optimizer_factory.py 中的优化器字典。当前仓库额外提供了运行时注册入口 register_optimizer_clsregister_optimizer_cls(key, optimizer_config_cls, use_legacy_optimizer)可以把自定义优化器类动态挂到LEGACY_OPTIMIZERS_CLS或NEW_OPTIMIZERS_CLS重复注册会抛ValueError。其 docstring 也明确提示仅注册类还不够用户仍需为它提供 oneof 配置 dataclass 才能被OptimizerFactory识别——这与文档三步流程中的第 2 步相互印证。学习率与 Warmup 调度学习率通过learning_rate字段必填配置warmup 通过warmup字段可选配置。文档列出的调度类型清单LR_CLS { stepwise: tf.keras.optimizers.schedules.PiecewiseConstantDecay, polynomial: tf.keras.optimizers.schedules.PolynomialDecay, exponential: tf.keras.optimizers.schedules.ExponentialDecay, cosine: tf.keras.experimental.CosineDecay, power: lr_schedule.DirectPowerDecay, } WARMUP_CLS { linear: lr_schedule.LinearWarmup, polynomial: lr_schedule.PolynomialWarmUp }此外还允许constant常数学习率类型从源码看optimizer_factory.py#L176-L179当type constant时build_learning_rate直接返回lr_config.learning_rate这个浮点值不构造 schedule 对象。当前仓库的 LR_CLS#L59-L69 已扩充为stepwise、polynomial、exponential、cosine、cosine_restarts、power、power_linear、power_with_offset、step_cosine_with_offset共九种其中前五种衰减类型均带WithOffset后缀是由 _make_offset_wrapper 动态生成的子类——行为等价于new_class_object(step) base_lr_class_object(step - offset)即给所有调度统一加上了一个offset步数偏移能力各配置 dataclass 中的offset字段默认 0 即为由此而来。各调度的配置字段与默认值定义在 learning_rate_config.py常用项摘录配置类关键字段默认值说明ConstantLrConfiglearning_rate: 0.1常数学习率StepwiseLrConfigboundaries、values均为必填、offset: 0values必须比boundaries多一个元素[0, b0] - v0、[b0, b1] - v1……ExponentialLrConfiginitial_learning_rate、decay_steps、decay_rate必填、staircase: False、offset: 0指数衰减可开启阶梯模式PolynomialLrConfigend_learning_rate: 0.0001、power: 1.0、cycle: False多项式衰减cycleTrue时可循环CosineLrConfiginitial_learning_rate、decay_steps必填、alpha: 0.0余弦衰减alpha是末端学习率占初始值比例CosineRestartsLrConfigfirst_decay_steps必填、t_mul: 2.0、m_mul: 1.0、alpha: 0.0余弦重启周期逐段拉长/重置DirectPowerLrConfigpower: -0.5遵循lr * step^power默认即平方根衰减PowerAndLinearDecayLrConfigtotal_decay_steps必填、power: -0.5、linear_decay_fraction: 0.1先 power 衰减末尾一段线性压到 0LinearWarmupConfigwarmup_learning_rate: 0、warmup_steps必填线性 warmup 起点学习率PolynomialWarmupConfigpower: 1、warmup_steps必填多项式 warmup学习率工作机制Warmup 与衰减如何拼接文档对学习率如何工作给出了三条明确的规则这是理解整套调度语义的关键学习率函数以step为输入返回该步的学习率值随训练推进学习率通常逐渐衰减warmup 调度常用于稳定训练过程它从一个较低的学习率起步逐步增大直至达到常规衰减调度的初始值拼接方式如下设 warmup 步数为warmup_steps总训练步数为train_steps步骤[0, warmup_steps)learning_rate warmup(step)步骤[warmup_steps, train_steps)learning_rate lr(step)warmup 的终点学习率不是用户指定的而是由衰减调度反推得到的即learning_rate(warmup_steps) warmup(warmup_steps)。并且 warmup不会把常规衰减整体推迟 warmup_steps 步而是替换掉这段区间的衰减。LinearWarmup 的实现与上述语义一一对应。构造时若主调度是 schedule 对象则执行self._final_warmup_lr after_warmup_lr_sched(warmup_steps)L131-L135——这正是终点值由衰减调度在warmup_steps处求值反推的源码体现。每一步的取值逻辑为linear_warmup_lr ( self._init_warmup_lr global_step / self._warmup_steps * (self._final_warmup_lr - self._init_warmup_lr)) lr tf.cond(global_step self._warmup_steps, lambda: linear_warmup_lr, lambda: after_warmup_lr) # 主调度按原始 step 求值不做偏移注意after_warmup_lr用的是未经偏移的原始step而非step - warmup_steps这从代码层面印证了替换而非推迟的设计warmup 结束后衰减曲线的位置与没有 warmup 时完全一致只是[0, warmup_steps)区间被线性爬升覆盖。同理 PolynomialWarmUp 中warmup_learning_rate initial_learning_rate * (step / warmup_steps)^power且对step0做了tf.math.maximum(step, 1.0)保护以避免除零产生 Inf。关于学习率日志的可见性文档还提醒学习率数值是按 summary_intervalTrainerConfig.summary_interval即两次 summary 写入之间的步数周期性记录的。如果warmup_steps小于summary_interval你在 summary 里就看不到 warmup 阶段的取值。以 imagenet_resnet50_tpu.yaml 为例summary_interval: 312而warmup_steps: 15605 个采样点warmup 曲线完整可见反之若把 warmup 设成 100 步则大概率整个 warmup 过程都不会出现在学习率曲线上——调参观察时应留意这一点。示例Cosine 衰减 线性 Warmup文档给出的另一个例子指定 decay_steps 为 20000 的余弦学习率衰减并对前 500 步做线性 warmupparams {learning_rate: {type: cosine, cosine: {decay_steps: 20000}}, warmup: {type: linear, linear: {warmup_steps: 500}}}按前文的拼接规则[0, 500)区间内学习率从warmup_learning_rate默认 0线性爬升到cosine(500)处反推出的值[500, 20000)则严格沿余弦曲线衰减曲线位置不因 warmup 而后移。实战视角视觉实验配置中的 optimizer_config在 TFM 视觉模块的实验配置里优化配置统一挂在trainer段的optimizer_config字段下。imagenet_resnet50_tpu.yaml#L27-L47 是一份可直接对照文档示例的完整实例trainer: train_steps: 28080 validation_steps: 13 validation_interval: 312 steps_per_loop: 312 summary_interval: 312 checkpoint_interval: 312 optimizer_config: optimizer: type: sgd sgd: momentum: 0.9 learning_rate: type: stepwise stepwise: boundaries: [9360, 18720, 24960] values: [1.6, 0.16, 0.016, 0.0016] warmup: type: linear linear: warmup_steps: 1560这份配置恰好演示了文档重要考量两点的实际落地train steps 与衰减边界的联动boundaries的 9360 / 18720 / 24960 分别是train_steps28080 的 1/3、2/3 和 5/7——衰减边界是相对总步数按比例设定的而不是拍脑袋的绝对值batch size 与学习率的联动该实验global_batch_size: 4096对应初始学习率 1.61/4096 规模的线性缩放结果warmup 步数 1560 约等于 10 个 epoch 内的一个固定爬升段。若把全局 batch size 减半或翻倍学习率各档位1.6 → 0.16 → 0.016 → 0.0016与训练步数都需要同步重估。optimizer_config这一字段在训练框架中的承接者是 TrainerConfig它与train_steps、summary_interval等调度类参数同处一个配置段保证训练节奏与优化节奏在一份配置中集中维护。在 Task 中定制优化器文档指出优化器与学习率的创建发生在 Task 内部如果任务需要不同的优化器/学习率调度器可以通过覆写类方法实现。对应源码是 Task.create_optimizerclassmethodclassmethod def create_optimizer(cls, optimizer_config: OptimizationConfig, runtime_config: Optional[RuntimeConfig] None, dp_config: Optional[DifferentialPrivacyConfig] None): gradient_transformers None if dp_config is not None: noise_stddev dp_config.clipping_norm * dp_config.noise_multiplier gradient_transformers [ functools.partial(ops.clip_l2_norm, l2_norm_clipdp_config.clipping_norm), functools.partial(ops.add_noise, noise_stddevnoise_stddev) ] opt_factory optimization.OptimizerFactory(optimizer_config) optimizer opt_factory.build_optimizer( opt_factory.build_learning_rate(), gradient_transformersgradient_transformers) if runtime_config: optimizer performance.configure_optimizer( optimizer, use_float16runtime_config.mixed_precision_dtype float16, loss_scaleruntime_config.loss_scale) return optimizer从源码结构看这个默认实现在构建优化器这一核心动作之外还叠加了三项框架级能力自定义任务覆写时可以按需保留或替换差分隐私梯度变换当传入dp_config时向build_optimizer注入clip_l2_normadd_noise两个 gradient transformer这正是 build_optimizer 签名 中gradient_transformers参数的用途且此时不应再设置clipvalue/clipnorm/global_clipnorm二者互斥混合精度配置若 runtime 配置了mixed_precision_dtype与loss_scale如 ResNet50 实验中的bfloat16会用performance.configure_optimizer为优化器套上 loss scale避免 float16 下溢出/下溢EMA 包装若OptimizationConfig.ema非空build_optimizer内部会自动把优化器包进 ExponentialMovingAverage要求use_legacy_optimizerTrue见 optimizer_factory.py#L248-L254。因此视觉/检测等项目里自定义优化器的标准姿势是在任务类中覆写create_optimizer这个类方法复用OptimizerFactory处理配置到实例的映射只在自己需要的位置如自定义梯度变换、额外的后处理postprocessor插入差异逻辑而不必从零实现调度逻辑。构建结果的正确性有 optimizer_factory_test.py 等测试文件覆盖可运行仓库测试验证你的覆写是否与默认行为一致。重要考量因素文档最后强调了两个调参时必须成对考虑的因素这里结合前面的配置实例再作总结Batch size改变 batch size 通常需要同比例缩放学习率数值和训练步数步数与样本吞吐相关。修改 batch size 后请同步调整learning_rate各档位数值与train_steps只改其一会破坏原有的等效训练预算Train steps训练总步数与decay_stepscosine、boundariesstepwise、total_decay_stepspower_linear等字段高度相关。文档明确警告只改其中一个不改另一个可能导致非预期的行为。以 imagenet_resnet50_tpu.yaml 为例若把train_steps从 28080 翻倍却不动boundaries: [9360, 18720, 24960]衰减将在总步数不到一半时就全部走完后半程训练将以最低档学习率空转——这就是文档所指的非预期行为。关键文件索引内容仓库路径本文对应的官方文档optimization.md优化器工厂构建调度与优化器、类型注册optimizer_factory.py优化器配置 dataclass含裁剪字段、默认值optimizer_config.py学习率/warmup 配置 dataclass含默认值learning_rate_config.py顶层 OptimizationConfigoptimizer/ema/learning_rate/warmup 聚合optimization_config.py调度实现LinearWarmup、PolynomialWarmUp、Power 系衰减、offset 包装lr_schedule.pyoneof 配置语义oneof.pyTask.create_optimizer 默认实现DP、混合精度base_task.pyTrainerConfigsummary_interval 等训练节奏参数config_definitions.py完整实验配置示例imagenet_resnet50_tpu.yaml【免费下载链接】modelsModels and examples built with TensorFlow项目地址: https://gitcode.com/GitHub_Trending/mode/models创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考