tf.tpu.experimental.embedding.TPUEmbeddingV0

API среднего уровня TPUEmbedding, работающий на TPU без ускорителя Embedding.

Просмотр псевдонимов

Псевдонимы совместимости для миграции

См. Руководство по миграции для получения более подробной информации.

tf.compat.v1.tpu.experimental.embedding.TPUEmbeddingV0

tf.tpu.experimental.embedding.TPUEmbeddingV0(
    feature_config: Union[tf.tpu.experimental.embedding.FeatureConfig, Iterable],
    optimizer: Optional[tpu_embedding_v2_utils._Optimizer]
)

Примечание: Этот API среднего уровня не предназначен для крупномасштабного поиска в таблицах встраивания. Таблицы встраивания будут дублироваться по устройствам, а не фрагментироваться по ним. Для крупномасштабного поиска в таблицах встраивания используйте класс tpu.experimental.embedding.TPUEmbedding. Этот класс — альтернативный способ поиска встраивания, когда TPU не поддерживает ни одну версию функции встраивания. Подробное объяснение см. в tpu.experimental.tpu_hardware_feature.embedding_feature.

Этот класс должен быть создан в TPUStrategy, иначе будет возбуждена ошибка RuntimeError.

strategy = tf.distribute.TPUStrategy(...)
with strategy.scope():
  embedding = tf.tpu.experimental.embedding.TPUEmbeddingV0(
      feature_config=feature_config,
      optimizer=tf.tpu.experimental.embedding.SGD(0.1))

При создании распределенного набора данных, который должен быть передан в операцию поиска, необходимо указать специальный параметр входа:

distributed_dataset = (
    strategy.distribute_datasets_from_function(
        dataset_fn=...,
        options=tf.distribute.InputOptions(
            experimental_fetch_to_device=False))
dataset_iterator = iter(distributed_dataset)

Ниже приведен пример шага обучения и оценки:

optimizer = tf.keras.optimizers.SGD(0.1)

@tf.function
def training_step(dataset_iterator, num_steps):
  def tpu_step(embedding_features):
    with tf.GradientTape() as tape:
      tape.watch(embedding.embedding_table.values())
      activation = embedding(embedding_features)
      model_output = model(activations)
      loss = ...  # some function of labels and model_output

    embedding_gradients = tape.gradient(loss,
                                        embedding.embedding_table.values())
    optimizer.apply_gradients(list(zip(gradients,
                              mid_level_api.embedding_tables.values())))
    # Insert your model gradient and optimizer application here

  for _ in tf.range(num_steps):
    strategy.run(tpu_step, args=(next(dataset_iterator), ))

@tf.function
def evalution_step(dataset_iterator, num_steps):
  def tpu_step(embedding_features):
    activations = embedding(embedding_features)
    model_output = model(activations)
    # Insert your evaluation code here.

  for _ in tf.range(num_steps):
    strategy.run(tpu_step, args=(next(dataset_iterator), ))

Примечание: Здесь используется оптимизатор Keras. Для обеспечения согласованности создания переменных слотов между оптимизаторами Keras и оптимизаторами встраивания необходимо передать аргумент slot_variable_creation_fn оптимизаторов встраивания с помощью функции Keras add_slot. Также обратите внимание, что имена слотов могут незначительно отличаться.

optimizer = tf.keras.optimizers.Adagrad(learning_rate=0.1)

def slot_variable_creation_fn(table, slot_names, slot_initializers):
    slots = {}
    for slot, initializer in zip(slot_names, slot_initializers):
      slots[slot] = optimizer.add_slot(table, slot, initializer)
    return slots

embedding_optimizer = tf.experimental.embedding.Adagrad(
    learning_rate=0.1,
    slot_variable_creation_fn=slot_variable_creation_fn)

# Use the embedding optimizer to create mid level api and keras optimizer to
# apply gradients.

Атрибуты
`embedding_tables`	Возвращает словарь таблиц встраивания, ключи которого — `TableConfig`.

Методы

`build`

Просмотреть исходный код

build()

Создает переменные и переменные слотов для встраивания TPU.

`embedding_lookup`

Просмотреть исходный код

embedding_lookup(
    features: Any, weights: Optional[Any] = None
) -> Any

Применяет поиск встраивания на TPU с использованием Tensorcore.

Обратите внимание, что все разреженные и фрагментированные тензоры будут преобразованы в плотные тензоры на ЦП, а затем переданы в TPU для поиска встраивания. Крупномасштабный поиск встраивания не поддерживается этим API, используйте API среднего уровня TPUEmbedding вместо этого.

Аргументы

features вложенная структура тензоров, разреженных тензоров или фрагментированных тензоров.

weights вложенная структура тензоров, разреженных тензоров или фрагментированных тензоров или None для отсутствия весов. Если не None, структура должна соответствовать структуре входных данных, но допускается наличие элементов None.

Аргументы
`features`	вложенная структура тензоров, разреженных тензоров или фрагментированных тензоров.
`weights`	вложенная структура тензоров, разреженных тензоров или фрагментированных тензоров или None для отсутствия весов. Если не None, структура должна соответствовать структуре входных данных, но допускается наличие элементов None.

Возвращает
Вложенная структура тензоров с той же структурой, что и входные данные.

`call`

Просмотреть исходный код

__call__(
    features: Any, weights: Optional[Any] = None
) -> Any

Вызов API среднего уровня для поиска встраивания.

© 2022 The TensorFlow Authors. All rights reserved.
Licensed under the Creative Commons Attribution License 4.0.
Code samples licensed under the Apache 2.0 License.
https://www.tensorflow.org/api_docs/python/tf/tpu/experimental/embedding/TPUEmbeddingV0