TensorFlow自定义优化器方法详解

本文深入剖析了在 TensorFlow 中自定义优化器的核心要点与实战陷阱，强调不能直接修改抽象基类 `tf.keras.optimizers.Optimizer` 的 `apply_gradients` 方法，而必须完整实现该方法及配套的 `_resource_apply_dense`/`_resource_apply_sparse`、`get_config`/`from_config`，同时严格遵循设备对齐、图模式兼容、状态管理与序列化规范；文章直击开发者常见误区——如遗漏关键方法、误用 Python 控制流、忽视稀疏梯度支持、错误处理学习率调度和 slot 变量放置等，以清晰原理+典型报错+正确范式的方式，帮你避开性能断崖与训练崩溃，真正写出健壮、高效、可保存可分布的自定义优化器。

为什么不能直接改 tf.keras.optimizers.Optimizer 的 apply_gradients

因为 tf.keras.optimizers.Optimizer 是抽象基类，apply_gradients 默认抛出 NotImplementedError。你继承它时，必须重写这个方法，否则运行时报错：NotImplementedError: apply_gradients() must be implemented in descendants。

常见错误是只重写了 _create_slots 或 _resource_apply_dense，却漏掉 apply_gradients —— 这个方法才是优化器真正“干活”的入口，Keras 训练循环（model.train_step）最终调用的就是它。

• apply_gradients 必须接收 grads_and_vars（梯度-变量对列表），返回一个 tf.Operation 或 tf.Tensor（通常用 tf.group 合并更新操作）
• 内部要手动调用 self._create_slots（如果用了动量、二阶矩等状态），否则第一次训练就报 LookupError: Slot does not exist
• 别在 apply_gradients 里写 Python 控制流（如 if 判断梯度是否为 None），要用 tf.cond 或 tf.where，否则图模式下会断图

怎么正确实现 _resource_apply_dense 和 _resource_apply_sparse

这两个方法才是真正更新单个变量的底层逻辑：_resource_apply_dense 处理稠密梯度（最常见），_resource_apply_sparse 处理稀疏梯度（如 embedding lookup 场景）。Keras 在 apply_gradients 内部按梯度类型分发调用它们。

容易踩的坑是只实现 _resource_apply_dense，结果遇到 SparseTensor 梯度时报错：TypeError: Expected dense grad, got sparse。即使你模型里没显式用 sparse op，某些 layer（如 tf.keras.layers.Embedding）默认产生 sparse 梯度。

• _resource_apply_dense 参数固定： grad（tf.Tensor）、var（tf.Variable）、apply_state（含学习率等共享状态，必须传给父类方法）
• _resource_apply_sparse 多一个 indices 参数，需用 tf.gather_nd 或 tf.scatter_nd 更新对应位置，不能直接加法
• 所有变量更新操作必须用 var.assign(...) 或 var.scatter_add(...)，不能用 var = var + ...（这只会新建张量，不更新原变量）

get_config 和 from_config 不写会怎样？

不写的话，模型保存（model.save）或分布式策略（tf.distribute.MirroredStrategy）下会失败，报错：ValueError: Optimizer state cannot be loaded when optimizer is not compiled 或更隐蔽的 AttributeError: 'dict' object has no attribute 'learning_rate'。

原因：Keras 序列化优化器时，先调 get_config() 取参数字典，再用 cls.from_config(config) 重建实例。如果你自定义优化器没实现它们，Keras 就只能存个空字典，反序列化时构造不出正确对象。

• get_config 返回 dict，键名必须和 __init__ 参数名一致（如 'learning_rate'、'beta_1'），值要是 JSON-serializable 类型（数字、字符串、布尔）
• from_config 必须调 cls(**config)，不能漏掉 ** 解包，否则参数传不进去
• 如果用了非标量参数（如 tf.Tensor 学习率），得在 get_config 里转成 Python 值（lr.numpy().item()），否则序列化失败

性能陷阱：别在 apply_gradients 里做重复计算

一个典型错误是在每次 apply_gradients 调用里重新计算全局学习率（比如带 warmup 的 schedule），导致图中插入大量冗余 op，训练变慢，且无法被 XLA 优化。

正确做法是把学习率逻辑移到 __init__ 或 build 阶段，用 tf.keras.optimizers.schedules.LearningRateSchedule 子类封装，然后在 apply_gradients 中通过 self._get_hyper('learning_rate') 获取当前值——这个方法会自动处理 schedule 的 step 更新。

• 所有与 step 相关的状态（如 self.iterations）必须用 self.add_weight 创建，并设 trainable=False，否则会被当作可训练参数参与梯度计算
• 避免在 apply_gradients 中调用 tf.print 或 tf.debugging.assert_*，这些 op 会强制同步，大幅拖慢 GPU 利用率
• 如果要加 custom clip（如 per-layer gradient norm），优先用 tf.clip_by_global_norm，而不是对每个变量单独 clip，前者一次算范数，后者多次访存

最常被忽略的是 slot 变量的 device placement：如果你在 CPU 上创建了 self._momentums，但模型在 GPU 上，每次更新都要跨设备拷贝，性能暴跌。务必在 _create_slots 里用 with tf.device(var.device): 对齐变量位置。

文中关于的知识介绍，希望对你的学习有所帮助！若是受益匪浅，那就动动鼠标收藏这篇《TensorFlow自定义优化器方法详解》文章吧，也可关注golang学习网公众号了解相关技术文章。