tf.tpu.experimental.embedding.TPUEmbedding

API среднего уровня TPUEmbedding.

tf.tpu.experimental.embedding.TPUEmbedding(
    feature_config: Any,
    optimizer: Optional[tpu_embedding_v2_utils._Optimizer],
    pipeline_execution_with_tensor_core: bool = False
)

Примечание: При создании внутри TPUStrategy этот класс можно создать только один раз за вызов tf.tpu.experimental.initialize_tpu_system. Если вы хотите повторно инициализировать движок встраивания, необходимо также повторно инициализировать TPU. При этом все переменные на TPU будут очищены, поэтому убедитесь, что вы выполнили контрольную точку до этого. Если требуется больше экземпляров класса, задайте аргумент initialize_tpu_embedding значение False.

Этот класс можно использовать для поддержки обучения больших встраиваний на TPU. При создании экземпляра этого класса необходимо указать полный набор таблиц и функций, которые вы ожидаете найти в этих таблицах. См. документацию по tf.tpu.experimental.embedding.TableConfig и tf.tpu.experimental.embedding.FeatureConfig для получения более подробной информации о полном наборе параметров. Здесь мы рассмотрим основное использование.

Примечание: несколько FeatureConfig объектов могут использовать один и тот же TableConfig объект, позволяя различным функциям совместно использовать одну и ту же таблицу:

table_config_one = tf.tpu.experimental.embedding.TableConfig(
    vocabulary_size=...,
    dim=...)
table_config_two = tf.tpu.experimental.embedding.TableConfig(
    vocabulary_size=...,
    dim=...)
feature_config = {
    'feature_one': tf.tpu.experimental.embedding.FeatureConfig(
        table=table_config_one),
    'feature_two': tf.tpu.experimental.embedding.FeatureConfig(
        table=table_config_one),
    'feature_three': tf.tpu.experimental.embedding.FeatureConfig(
        table=table_config_two)}

Существует два режима использования класса TPUEmbedding. Это зависит от того, был ли класс создан внутри области TPUStrategy или нет.

Внутри TPUStrategy, мы разрешаем доступ к методу enqueue, dequeue и apply_gradients. Ниже мы покажем примеры использования этих методов для обучения и оценки вашей модели. В режиме CPU мы имеем доступ только к свойству embedding_tables, которое позволяет получить доступ к таблицам встраивания, чтобы вы могли использовать их для выполнения оценки/предсказания модели на CPU.

Сначала рассмотрим режим TPUStrategy. Начальная настройка выглядит так:

strategy = tf.distribute.TPUStrategy(...)
with strategy.scope():
  embedding = tf.tpu.experimental.embedding.TPUEmbedding(
      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_prefetch_to_device=False))
dataset_iterator = iter(distributed_dataset)

Примечание: Все пакеты, передаваемые слою, должны иметь одинаковый размер пакета для каждого входного элемента. Более того, после вызова слоя с одним размером пакета все последующие вызовы должны использовать тот же размер пакета. В случае, если размер пакета не может быть автоматически определён методом очереди, необходимо вызвать метод build со значением размера пакета для инициализации слоя.

Чтобы использовать этот API на TPU, следует использовать пользовательский цикл обучения. Ниже приведён пример шага обучения и оценки:

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

    embedding_gradients = tape.gradient(loss, activations)
    embedding.apply_gradients(embedding_gradients)
    # Insert your model gradient and optimizer application here

  for _ in tf.range(num_steps):
    embedding_features, tpu_features = next(dataset_iterator)
    embedding.enqueue(embedding_features, training=True)
    strategy.run(tpu_step, args=(embedding_features, ))

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

  for _ in tf.range(num_steps):
    embedding_features, tpu_features = next(dataset_iterator)
    embedding.enqueue(embedding_features, training=False)
    strategy.run(tpu_step, args=(embedding_features, ))

Примечание: Вызовы enqueue имеют training установленное в True при использовании embedding.apply_gradients и установленное в False при отсутствии embedding.apply_gradients в функции. Если вы не придерживаетесь этой схемы, может возникнуть ошибка или TPU может зависнуть.

В приведённых выше примерах предполагается, что у пользователя есть набор данных, возвращающий кортеж, где первый элемент кортежа соответствует структуре того, что было передано в качестве аргумента feature_config инициализатору объекта. Также мы используем tf.range для получения tf.while_loop с целью повышения производительности.

При сохранении модели в контрольную точку необходимо включить объект tf.tpu.experimental.embedding.TPUEmbedding в контрольную точку. Это отслеживаемый объект, и сохранение его позволит сохранить таблицы встраивания и их переменные слота оптимизатора:

checkpoint = tf.train.Checkpoint(model=model, embedding=embedding)
checkpoint.save(...)

На CPU можно использовать только свойство embedding_table. Это позволит восстановить контрольную точку в объекте и получить доступ к переменным таблицы:

model = model_fn(...)
embedding = tf.tpu.experimental.embedding.TPUEmbedding(
    feature_config=feature_config,
    batch_size=1024,
    optimizer=tf.tpu.experimental.embedding.SGD(0.1))
checkpoint = tf.train.Checkpoint(model=model, embedding=embedding)
checkpoint.restore(...)

tables = embedding.embedding_tables

Теперь вы можете использовать таблицу в функциях, таких как tf.nn.embedding_lookup, для выполнения поиска встраивания и передачи в вашу модель.

Методы

apply_gradients

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

apply_gradients(
    gradients,
    name: Text = None
)

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

Если градиент None передан в любом положении вложенной структуры, то для этой функции применяется обновление градиента с нулевым градиентом. Для оптимизаторов, таких как SGD или Adagrad, это то же самое, что и не применять обновление вообще. Для ленивого Adam и других оптимизаторов с применением по частям, учитывайте влияние применения нулевого градиента.

strategy = tf.distribute.TPUStrategy(...)
with strategy.scope():
  embedding = tf.tpu.experimental.embedding.TPUEmbedding(...)

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

@tf.function
def training_step():
  def tpu_step(tpu_features):
    with tf.GradientTape() as tape:
      activations = embedding.dequeue()
      tape.watch(activations)

      loss = ... #  some computation involving activations

    embedding_gradients = tape.gradient(loss, activations)
    embedding.apply_gradients(embedding_gradients)

  embedding_features, tpu_features = next(dataset_iterator)
  embedding.enqueue(embedding_features, training=True)
  strategy.run(tpu_step, args=(embedding_features, ))

training_step()

build

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

build(
    per_replica_batch_size: Optional[int] = None
)

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

Этот метод создаёт базовые переменные (включая переменные слота). Если создан внутри TPUStrategy, это также инициализирует TPU для встраивания.

Эта функция автоматически вызывается enqueue, которая попытается автоматически определить размер вашего пакета. Если это не удастся, необходимо вручную вызвать этот метод перед вызовом enqueue.

dequeue

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

dequeue(
    name: Text = None
)

Получение результатов встраивания.

Возвращает вложенную структуру объектов tf.Tensor, соответствующую структуре аргумента feature_config класса TPUEmbedding. Форма выходных тензоров имеет вид (batch_size, dim), где batch_size — размер пакета на ядро, dim — размерность соответствующей TableConfig. Если у соответствующей FeatureConfig функции max_sequence_length больше 0, вывод будет представлять собой последовательность формы (batch_size, max_sequence_length, dim).

strategy = tf.distribute.TPUStrategy(...)
with strategy.scope():
  embedding = tf.tpu.experimental.embedding.TPUEmbedding(...)

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

@tf.function
def training_step():
  def tpu_step(tpu_features):
    with tf.GradientTape() as tape:
      activations = embedding.dequeue()
      tape.watch(activations)

      loss = ... #  some computation involving activations

    embedding_gradients = tape.gradient(loss, activations)
    embedding.apply_gradients(embedding_gradients)

  embedding_features, tpu_features = next(dataset_iterator)
  embedding.enqueue(embedding_features, training=True)
  strategy.run(tpu_step, args=(embedding_features, ))

training_step()

переданный в этот экземпляр объекта TPUEmbedding.

enqueue

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

enqueue(
    features,
    weights=None,
    training: bool = True,
    name: Optional[Text] = None
)

В очереди id тензоров для поиска вставленных значений.

Эта функция добавляет структуру признаков в очередь для поиска в таблицах вставленных значений. Мы ожидаем, что размер пакета каждого из тензоров в features соответствует размеру пакета на ядро. Это произойдёт автоматически, если ваш входной набор данных сгруппирован в глобальный размер пакета и вы используете tf.distribute.TPUStrategy's experimental_distribute_dataset или если вы используете distribute_datasets_from_function и группируете в размер пакета на ядро, вычисленный контекстом, переданным вашей функции ввода.

strategy = tf.distribute.TPUStrategy(...)
with strategy.scope():
  embedding = tf.tpu.experimental.embedding.TPUEmbedding(...)

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

@tf.function
def training_step():
  def tpu_step(tpu_features):
    with tf.GradientTape() as tape:
      activations = embedding.dequeue()
      tape.watch(activations)

      loss = ... #  some computation involving activations

    embedding_gradients = tape.gradient(loss, activations)
    embedding.apply_gradients(embedding_gradients)

  embedding_features, tpu_features = next(dataset_iterator)
  embedding.enqueue(embedding_features, training=True)
  strategy.run(tpu_step, args=(embedding_features,))

training_step()

Примечание: Вы должны указать training=True при использовании embedding.apply_gradients как указано выше и training=False при его отсутствии embedding.apply_gradients (например, для замороженных вставленных значений или при выполнении оценки).