Пользовательский оптимизатор TensorFlow Keras

Question

Предположим, я хочу написать собственный класс оптимизатора, который соответствует 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заглушка метода. Точно нетresource_apply_sparse или create_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: _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_slots, но база tf.keras.optimizers.Optimizer класс не имеет _create_slotsметод. Итак, похоже, что_create_slotsметод должен быть определен в подклассе оптимизатора, если этот подкласс не отменяетapply_gradients.

---

Вопросы

Как правильно создать подкласс 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.

Answer 1

Обновление: TF2.2 заставил меня очистить все реализации - теперь их можно использовать в качестве справочника для лучших практик TF. Также добавлен раздел ниже по_get_hyper vs. _set_hyper.

---

Я реализовал Keras AdamW во всех основных версиях TF и ​​Keras - приглашаю вас изучить optimizers_v2.py. Несколько моментов:

• Вы должны унаследовать OptimizerV2, который вы на самом деле связали; это последний и текущий базовый класс дляtf.keras оптимизаторы
• Вы правы в (1) - это ошибка документации; методы являются частными, поскольку они не предназначены для непосредственного использования пользователем.
• apply_gradients(или любой другой метод) отменяется только в том случае, если по умолчанию не выполняется то, что необходимо для данного оптимизатора; в вашем связанном примере это просто однострочный аддон к исходному
• "Итак, похоже, что _create_slots метод должен быть определен в подклассе оптимизатора, если этот подкласс не отменяет apply_gradients" - эти двое не связаны между собой; это случайно.

---

• В чем разница между _resource_apply_dense а также _resource_apply_sparse?

Последний имеет дело с разреженными слоями - например,Embedding- и прежний со всем остальным; пример.

• Когда мне следует использовать _create_slots()?

При определении обучаемого tf.Variables; пример: моменты весов первого и второго порядка (например, Адам). Оно использует add_slot().

---

_get_hyper vs. _set_hyper: они позволяют устанавливать и получать литералы Python (int, strи т. д.), вызываемые объекты и тензоры. Они существуют в основном для удобства: все, что установлено через_set_hyper можно получить через _get_hyper, избегая повторения шаблонного кода. Я посвятил этому вопрос и ответ.

Answer 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()

2. Я не знаю о 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.