tf.tpu.experimental.embedding.TPUEmbedding

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

tf.tpu.experimental.embedding.TPUEmbedding(
    feature_config, batch_size, optimizer,
    pipeline_execution_with_tensor_core=False, initialize_tpu_embedding=True
)

Примечание: При создании экземпляра в рамках 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,
      batch_size=1024,
      optimizer=tf.tpu.experimental.embedding.SGD(0.1))

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

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

Для использования этого 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=None
)

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

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

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

distributed_dataset = (
    strategy.experimental_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()

dequeue

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

dequeue(
    name=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.experimental_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=True, name=None
)

Загрузка тензоров id для поиска встраивания.

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

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

distributed_dataset = (
    strategy.experimental_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 (например, для замороженных встраиваний или при оценке).