tf.tpu.experimental.embedding.TPUEmbeddingV0

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

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.

Методы

build

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

build()

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

embedding_lookup

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

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

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

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

__call__

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

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

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