Optimizador personalizado TensorFlow Keras

30

Supongamos que quiero escribir una clase optimizadora personalizada que se ajuste a la tf.kerasAPI (usando la versión TensorFlow> = 2.0). Estoy confundido acerca de la forma documentada de hacer esto frente a lo que se hace en las implementaciones.

La documentación para los tf.keras.optimizers.Optimizer estados ,

  ### 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)

Sin embargo, la tf.keras.optimizers.Optimizerimplementación actual no define un resource_apply_densemétodo, pero define un _resource_apply_densecódigo auxiliar de aspecto privado . Del mismo modo, no hay resource_apply_sparseo create_slotsmétodos, pero hay un _resource_apply_sparsecódigo auxiliar de método y una _create_slotsllamada a método .

En oficiales tf.keras.optimizers.Optimizersubclases (utilizando tf.keras.optimizers.Adamcomo ejemplo), hay _resource_apply_dense, _resource_apply_sparsey _create_slotsmétodos, y no hay tales métodos sin el carácter de subrayado líder.

Existen métodos similares líder-subrayado en ligeramente menos oficiales tf.keras.optimizers.Optimizersubclases (por ejemplo, tfa.optimizers.MovingAveragedesde TensorFlow Addons: _resource_apply_dense, _resource_apply_sparse, _create_slots).

Otro punto de confusión para mí es que algunos de los optimizadores de complementos de TensorFlow también anulan el apply_gradientsmétodo (por ejemplo, tfa.optimizers.MovingAverage), mientras que los tf.keras.optimizersoptimizadores no.

Además, noté que el apply_gradientsmétodo de tf.keras.optimizers.Optimizermétodo llama_create_slots , pero la tf.keras.optimizers.Optimizerclase base no tiene un _create_slotsmétodo. Entonces, parece que un _create_slotsmétodo debe definirse en una subclase de optimizador si esa subclase no anula apply_gradients.


Preguntas

¿Cuál es la forma correcta de subclase a tf.keras.optimizers.Optimizer? Específicamente,

  1. ¿La tf.keras.optimizers.Optimizerdocumentación enumerada en la parte superior simplemente significa anular las versiones de guión bajo de los métodos que mencionan (por ejemplo, en _resource_apply_denselugar de resource_apply_dense)? Si es así, ¿hay alguna garantía API sobre estos métodos de aspecto privado que no cambian su comportamiento en futuras versiones de TensorFlow? ¿Cuáles son las firmas de estos métodos?
  2. ¿Cuándo se anularía apply_gradientsademás de los _apply_resource_[dense|sparse]métodos?

Editar. Problema abierto en GitHub: # 36449

Artem Mavrin
fuente
1
Esto puede ser algo para informar como un problema de documentación a los desarrolladores. Definitivamente, parece que esos métodos para anular deberían incluir el guión bajo inicial en la documentación, pero en cualquier caso, como usted dice, no hay información sobre su firma y el propósito exacto. También puede ser que los nombres de métodos sin guión bajo (y documentados) estén planeados para agregarse (como con get_config), pero que aún no deberían aparecer en la documentación pública .
jdehesa
Para las firmas, siempre puede mirar la declaración de _resource_apply_denseo _resource_apply_sparse, y ver su uso en optimizadores implementados. Si bien puede no ser, creo, una API pública con garantías de estabilidad, diría que es bastante seguro usarlas. Simplemente deberían proporcionar una mejor orientación en este aspecto.
jdehesa
Estoy de acuerdo en que este es un problema de documentación con TensorFlow. ¿Creó un problema para esto en el repositorio de tf Github? Si es así, ¿podría compartir el enlace aquí?
jpgard

Respuestas:

3

Implementé Keras AdamW en todas las versiones principales de TF & Keras. Los invito a examinar optimizers_v2.py . Varios puntos:

  • Deberías heredar OptimizerV2, que en realidad es lo que vinculaste; Es la clase base más reciente y actual para tf.kerasoptimizadores
  • Tiene razón en (1): este es un error de documentación; Los métodos son privados, ya que no están destinados a ser utilizados directamente por el usuario.
  • apply_gradients(o cualquier otro método) solo se anula si el valor predeterminado no logra lo que se necesita para un optimizador dado; en su ejemplo vinculado, es solo un complemento de una línea al original
  • "Entonces, parece que un _create_slotsmétodo debe definirse en una subclase de optimizador si esa subclase no anula apply_gradients" - los dos no están relacionados; Es coincidencia.

  • ¿Cuál es la diferencia entre _resource_apply_densey _resource_apply_sparse?

Último trata con capas dispersas , por ejemplo Embedding, y anterior con todo lo demás; ejemplo .

  • ¿Cuándo debo usar _create_slots()?

Al definir s entrenables tf.Variable ; ejemplo: momentos de primer y segundo orden de ponderaciones (por ejemplo, Adam). Lo utiliza add_slot().

Más o menos, cuando no se usa _create_slots(); es como establecer atributos de clase, pero con pasos de preprocesamiento adicionales para garantizar la corrección en el uso. Así Python int, float, tf.Tensor, tf.Variable, y otros. (Debería haberlo usado más en Keras AdamW).


Nota : aunque mis optimizadores vinculados funcionan correctamente y son casi tan rápidos como los originales, el código sigue las mejores prácticas de TensorFlow y aún puede ser más rápido; No lo recomiendo como la "referencia ideal". Por ejemplo, algunos objetos de Python (por ejemplo int) deberían ser tensores; eta_tse define como a tf.Variable, pero se anula inmediatamente como a tf.Tensoren los _applymétodos. No necesariamente es un gran problema, simplemente no he tenido tiempo para renovar.

OverLordGoldDragon
fuente
2
  1. Sí, esto parece ser un error de documentación. Los nombres de subrayado anteriores son los métodos correctos para anular. Relacionado está el Optimizador que no es de Keras, que tiene todo esto definido, pero no implementado en la clase base 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. No lo se apply_dense. Por un lado, si lo anula, el código menciona que una estrategia de distribución por réplica podría ser "peligrosa"
    # 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.
Tyler
fuente