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 не поддерживает ни одну версию функции embedding. Подробное объяснение см. в разделе 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 оптимизаторов эмбеддингов необходимо передать с функцией add_slot из Keras. Также обратите внимание, что имена слотов могут немного различаться.

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.

Обратите внимание, что все разреженные и ragged-тензоры будут преобразованы в плотные тензоры на CPU, а затем переданы на TPU для поиска эмбеддингов. Этот API не поддерживает поиск в больших таблицах эмбеддингов; вместо него используйте API TPUEmbedding среднего уровня.

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

Возвращает
Вложенную структуру тензоров с той же структурой, что и у inputs.

`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