tf.tpu.experimental.embedding.CustomOptimizer

Optimization parameters for custom optimizer for TPU embeddings.

tf.tpu.experimental.embedding.CustomOptimizer(
    custom_computation: tf.types.experimental.PolymorphicFunction,
    learning_rate: Union[float, Callable[[], float]] = 0.01,
    slot_names: Optional[List[str]] = None,
    slot_initializers: Optional[List[init_ops_v2.Initializer]] = None,
    hyperparameters: Optional[List[Union[float, Callable[[], float]]]] = None,
    slot_dims: Optional[List[int]] = None
) -> Any

This optimizer gives the user the ability to define a custom optimizer for running embedding lookups on TPU v5p.

The custom computation should be a function which takes gradient, embedding table, a list of slot variables, learning_rate and a list of hyperparameters. The function should perform the gradient update on the embedding_table + slot_variables and return the updated embedding_table and slot_variables. e.g.

@tf.function
def sgd_optimizer_computation(
    gradient,
    embedding_table,
    slot_variables,
    learning_rate,
    hyperparameters,
):
  del slot_variables, hyperparameters
  return embedding_table - gradient * learning_rate

Above is a simple example of a sgd optimizer. You can also define a more complex optimizer which updates multiple tables and slot variables.

@def_function.function
def adagrad_optimizer_computation(
    gradient,
    embedding_table,
    slot_variables,
    learning_rate,
    hyperparameters,
):
  del hyperparameters
  accumulator = slot_variables[0]
  new_accumulator = accumulator + gradient * gradient
  updated_embedding_table = (
      embedding_table
      - learning_rate * gradient / math_ops.sqrt(new_accumulator)
  )
  return (updated_embedding_table, new_accumulator)

The custom computation is defined as a per-row update function and it will be auto scaled for the entire table (slot variables).

Pass this to tf.tpu.experimental.embedding.TPUEmbeddingV2 via the optimizer argument to set the global optimizer and its parameters:

optimizer = tf.tpu.experimental.embedding.CustomOptimizer(
      custom_computation=adagrad_optimizer_computation,
      learning_rate=1.0,
      slot_names=['accumulators'],
      slot_initializers=[
          tf.constant_initializer(0.1, support_partition=True)
      ],
  )
embedding = tf.tpu.experimental.embedding.TPUEmbeddingV2(
    ...
    optimizer=optimizer)

This can also be used in a tf.tpu.experimental.embedding.TableConfig as the optimizer parameter to set a table specific optimizer. This will override the optimizer and parameters for global embedding optimizer defined above:

table_one = tf.tpu.experimental.embedding.TableConfig(
    vocabulary_size=...,
    dim=...,
    optimizer=optimizer)
table_two = tf.tpu.experimental.embedding.TableConfig(
    vocabulary_size=...,
    dim=...)

feature_config = (
    tf.tpu.experimental.embedding.FeatureConfig(
        table=table_one),
    tf.tpu.experimental.embedding.FeatureConfig(
        table=table_two))

embedding = tf.tpu.experimental.embedding.TPUEmbedding(
    feature_config=feature_config,
    batch_size=...
    optimizer=optimizer)

In this example, the optimizer of the first table will be the one specified in the table config. The second table will use the optimizer specified in the TPUEmbedding argument.

custom_computation

hyperparameters

slot_dims Optional per-slot dimensions metadata for variable creation.

Each entry corresponds to the same-index item in slot_names. This is consumed by the embedding layer when creating slot variables.

Methods

__eq__

View source

__eq__(
    other: Any
) -> Union[Any, bool]

Return self==value.