自定义 TensorFlow Keras 优化器

Posted 2023-02-16

【中文标题】自定义 TensorFlow Keras 优化器

【英文标题】：Custom TensorFlow Keras optimizer

【发布时间】：2020-03-05 10:58:42

【问题描述】：

假设我想编写一个符合tf.keras API 的自定义优化器类（使用 TensorFlow 版本>=2.0）。我对记录的执行此操作的方法与实现中所做的操作感到困惑。

tf.keras.optimizers.Optimizerstates 的文档，

  ### Write a customized optimizer.
  If you intend to create your own optimization algorithm, simply inherit from
  this class and override the following methods:

    - resource_apply_dense (update variable given gradient tensor is dense)
    - resource_apply_sparse (update variable given gradient tensor is sparse)
    - create_slots (if your optimizer algorithm requires additional variables)

然而，当前的tf.keras.optimizers.Optimizer 实现没有定义resource_apply_dense 方法，但它确实定义了一个看起来私有的_resource_apply_dense method stub。同样，没有resource_apply_sparse 或create_slots 方法，但有_resource_apply_sparse method stub 和_create_slots method call。

在官方的tf.keras.optimizers.Optimizer子类中（以tf.keras.optimizers.Adam为例），有_resource_apply_dense、_resource_apply_sparse和_create_slots方法，没有前导下划线的就没有这些方法。

在稍微不那么官方的 tf.keras.optimizers.Optimizer 子类中也有类似的前导下划线方法（例如，来自 TensorFlow Addons 的 tfa.optimizers.MovingAverage：_resource_apply_dense、_resource_apply_sparse、_create_slots）。

另一个让我困惑的地方是，一些 TensorFlow Addons 优化器也覆盖了 apply_gradients 方法（例如，tfa.optimizers.MovingAverage），而 tf.keras.optimizers 优化器没有。

此外，我注意到tf.keras.optimizers.Optimizer 方法calls _create_slots 的apply_gradients 方法，但是基础tf.keras.optimizers.Optimizer 类没有_create_slots 方法。 因此，如果子类不覆盖apply_gradients，则似乎必须在优化器子类中定义_create_slots 方法。

---

问题

子类化tf.keras.optimizers.Optimizer 的正确方法是什么？具体来说，

顶部列出的tf.keras.optimizers.Optimizer 文档是否仅仅意味着覆盖他们提到的方法的前导下划线版本（例如，_resource_apply_dense 而不是resource_apply_dense）？如果是这样，是否有任何 API 保证这些看起来私有的方法在未来版本的 TensorFlow 中不会改变它们的行为？这些方法的签名是什么？ 除了_apply_resource_[dense|sparse] 方法之外，何时会覆盖apply_gradients？

---

编辑。 在 GitHub 上打开的问题：#36449

【问题讨论】：

这可能会作为文档问题报告给开发人员。看起来那些要覆盖的方法肯定应该在文档中包含初始下划线，但无论如何，就像你说的那样，没有关于它们的签名和确切用途的信息。也可能是计划添加不带下划线（并记录在案）的方法名称（如get_config），但它们不应该出现在public documentation中。

对于签名，您可以随时查看_resource_apply_dense 或_resource_apply_sparse 的声明，并查看它们在实现的优化器中的用法。虽然我认为它可能不是具有稳定性保证的公共 API，但我认为使用它们非常安全。他们只是应该在这方面提供更好的指导。

我同意这是 TensorFlow 的文档问题。您是否在 tf Github 存储库中为此创建了问题？如果是这样，你能在这里分享链接吗？

【参考方案1】：

更新：TF2.2 强制我清理所有实现 - 所以现在它们可以用作 TF 最佳实践的参考。还在下面添加了一个关于_get_hyper 与_set_hyper 的部分。

---

我已经在所有主要的 TF 和 Keras 版本中实现了Keras AdamW - 我邀请您检查optimizers_v2.py。几点：

你应该继承OptimizerV2，这实际上是你链接的；它是 tf.keras 优化器的最新和当前基类 您在 (1) 中是正确的 - 这是一个文档错误；这些方法是私有的，因为它们不打算由用户直接使用。 apply_gradients（或任何其他方法）仅在默认值无法满足给定优化器所需的情况下才会被覆盖；在您的链接示例中，它只是原版的单行插件 “因此，如果子类不覆盖 apply_gradients，那么似乎必须在优化器子类中定义 _create_slots 方法” - 这两者是不相关的；纯属巧合。

---

_resource_apply_dense 和 _resource_apply_sparse 有什么区别？

Latter 处理 sparse 层 - 例如Embedding - 和以前的一切； example.

什么时候应该使用_create_slots()？

定义可训练时tf.Variables；示例：权重的一阶和二阶矩（例如 Adam）。它使用add_slot()。

---

_get_hyper 与 _set_hyper：它们支持设置和获取 Python 文字（int、str 等）、可调用对象和张量。它们的存在主要是为了方便：通过_set_hyper 设置的任何内容都可以通过_get_hyper 检索，避免重复样板代码。我专门对它进行了问答here。

【参考方案2】：

是的，这看起来是一个文档错误。前面的下划线名称是正确的覆盖方法。相关的是非 Keras 优化器，它已经定义了所有这些，但没有在基类 https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/training/optimizer.py 中实现

  def _create_slots(self, var_list):
    """Create all slots needed by the variables.
    Args:
      var_list: A list of `Variable` objects.
    """
    # No slots needed by default
    pass

  def _resource_apply_dense(self, grad, handle):
    """Add ops to apply dense gradients to the variable `handle`.
    Args:
      grad: a `Tensor` representing the gradient.
      handle: a `Tensor` of dtype `resource` which points to the variable
       to be updated.
    Returns:
      An `Operation` which updates the value of the variable.
    """
    raise NotImplementedError()

  def _resource_apply_sparse(self, grad, handle, indices):
    """Add ops to apply sparse gradients to the variable `handle`.
    Similar to `_apply_sparse`, the `indices` argument to this method has been
    de-duplicated. Optimizers which deal correctly with non-unique indices may
    instead override `_resource_apply_sparse_duplicate_indices` to avoid this
    overhead.
    Args:
      grad: a `Tensor` representing the gradient for the affected indices.
      handle: a `Tensor` of dtype `resource` which points to the variable
       to be updated.
      indices: a `Tensor` of integral type representing the indices for
       which the gradient is nonzero. Indices are unique.
    Returns:
      An `Operation` which updates the value of the variable.
    """
    raise NotImplementedError()

我不知道apply_dense。一方面，如果您确实覆盖它，代码会提到每个副本的 DistributionStrategy 可能是“危险的”

    # TODO(isaprykin): When using a DistributionStrategy, and when an
    # optimizer is created in each replica, it might be dangerous to
    # rely on some Optimizer methods.  When such methods are called on a
    # per-replica optimizer, an exception needs to be thrown.  We do
    # allow creation per-replica optimizers however, because the
    # compute_gradients()->apply_gradients() sequence is safe.