36

假设我想编写一个符合tf.kerasAPI 的自定义优化器类（使用 TensorFlow 版本>=2.0）。我对记录在案的方法与在实现中所做的事情感到困惑。

tf.keras.optimizers.Optimizer 状态的文档，

  ### 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方法 stub。同样，没有resource_apply_sparseorcreate_slots方法，但有一个_resource_apply_sparse方法存根和一个_create_slots方法调用。

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

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

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

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

问题

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

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

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

4

2 回答 2

6

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

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

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

后者处理稀疏层——例如Embedding——而前者处理其他一切；例子。

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

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

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

于 2020-02-20T01:14:52.190 回答
2
  1. 是的，这看起来是一个文档错误。前面的下划线名称是正确的覆盖方法。相关的是非 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()
  1. 我不知道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.
于 2020-02-19T21:32:40.847 回答