自定义TensorFlow Keras优化器

36
假设我想编写一个符合 TensorFlow 版本>=2.0 的自定义优化器类,符合 tf.keras API。我对文档中所述的方法与实现中的方法感到困惑。 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方法存根。同样,也没有resource_apply_sparsecreate_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子类中有类似的以下划线开头的方法 (例如,从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.Optimizerapply_gradients方法调用_create_slots,但基类tf.keras.optimizers.Optimizer没有_create_slots方法。 因此,如果子类没有覆盖apply_gradients,则似乎必须在优化器子类中定义_create_slots方法。

问题

如何正确地对tf.keras.optimizers.Optimizer进行子类化?具体而言,

  1. tf.keras.optimizers.Optimizer文档中所列举的是否只是指覆盖提到的方法(例如,使用_resource_apply_dense代替resource_apply_dense)?如果是这样,那么这些看起来像私有方法的API保证不会在未来版本的TensorFlow中更改其行为吗?这些方法的签名是什么?
  2. 除了使用_apply_resource_[dense|sparse]方法之外,何时需要覆盖apply_gradients

编辑。 在 GitHub 上打开了问题:#36449


1
这可能是一个需要向开发人员报告文档问题的事情。它看起来那些应该被覆盖的方法在文档中应该包含初始下划线,但无论如何,就像你所说的,没有关于它们签名和确切目的的信息。也可能是计划添加没有下划线(并记录在文档中)的方法名称(例如get_config),但是它们不应该出现在公共文档中。 - jdehesa
对于方法签名,您可以始终查看 _resource_apply_dense_resource_apply_sparse 的声明,并查看它们在实现的优化器中的用法。虽然可能不是公共API且稳定性有保证,但我认为它们是相当安全使用的。它们只应该在这方面提供更好的指导。 - jdehesa
我同意这是TensorFlow的文档问题。您是否在tf Github存储库中创建了此问题?如果是,能否在此处分享链接? - jpgard
2个回答

7

更新:TF2.2 强制我整理所有实现代码 - 现在它们可以用作 TF 最佳实践的参考。下面还添加了一个关于 _get_hyper_set_hyper 的部分说明。


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

  • 应该继承OptimizerV2,这实际上是您链接的内容;它是 tf.keras 优化器的最新和当前基类。
  • 您在第一点中是正确的 - 这是文档错误;这些方法是私有的,因为它们不打算由用户直接使用。
  • apply_gradients(或任何其他方法)仅在默认情况下无法完成给定优化器所需的任务时才被重写;在您链接的示例中,它只是原始代码的一个一行附加项。
  • "因此,在优化器子类中定义 _create_slots 方法是必需的,如果该子类没有覆盖 apply_gradients" - 这两者是不相关的;这是巧合。

  • _resource_apply_dense_resource_apply_sparse 之间有什么区别?

后者处理稀疏层 - 如 Embedding - 前者处理其他所有内容;示例

  • 何时应使用 _create_slots()

当定义可训练的 tf.Variable 时;例如:权重的一阶和二阶矩(如 Adam)。它使用add_slot()


_get_hyper_set_hyper:它们可以设置和获取 Python 文字(intstr 等)、可调用函数和张量。它们主要是为了方便:通过 _set_hyper 设置的任何内容都可以通过 _get_hyper 获取,避免了重复编写模板代码。我在这里专门讲解了它的问题和答案


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.

网页内容由stack overflow 提供, 点击上面的
可以查看英文原文,