tf.tpu.experimental.embedding.TPUEmbedding

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

tf.tpu.experimental.embedding.TPUEmbedding(
    feature_config: Union[tf.tpu.experimental.embedding.FeatureConfig, Iterable],
    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))

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

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

Разные входные функции могут иметь разные формы. Для плотных и разреженных тензоров поддерживается ранг 2 и выше. Для тензоров с неравномерным размером, хотя поддерживается только ранг 2, вы можете указать форму выходных данных для ранга 2 и выше. Указанная в FeatureConfig форма выходных данных имеет наивысший приоритет. Форма входных данных, переданная в метод build, имеет второй приоритет, а формы входных данных, автоматически определенные из входной функции, имеют наименьший приоритет. Последние два будут преобразованы в формы выходных данных путем исключения последнего измерения. Если у объектов с низким приоритетом есть формы выходных данных, которые не совпадают с предыдущими, будет поднято исключение ValueError. Только если у предыдущих объектов нет определенной формы выходных данных, последующие могут ее переопределить.

Примечание: Все пакеты, передаваемые в слой, могут иметь разные формы входных данных. Но эти формы должны совпадать с формами выходных данных, заданными либо FeatureConfig, либо методом build, за исключением тензоров с неравномерным размером. Разрешаются только тензоры с неравномерным размером 2D с формой выходных данных, заданной для более высоких измерений, при условии, что общее количество элементов совпадает. Все последующие вызовы должны иметь одинаковые формы входных данных. В случае, если формы входных данных не могут быть автоматически определены методом enqueue, необходимо вызвать метод build с формами входных данных или указать формы выходных данных в FeatureConfig для инициализации слоя.

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

@tf.function
def evaluation_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=(tpu_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,
    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: Optional[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_fetch_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=(tpu_features, ))

training_step()

build

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

build(
    per_replica_input_shapes=None, per_replica_batch_size=None
)

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

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

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

dequeue

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

dequeue(
    name: Optional[Text] = None
)

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

Возвращает вложенную структуру объектов tf.Tensor, соответствующую структуре аргумента feature_config класса TPUEmbedding. Форма выходных тензоров — (*output_shape, dim), dim — размер соответствующего TableConfig. Для output_shape есть три места, где она может быть установлена.

1. Конфигурация признаков, предоставленная в функции init.
2. Per_replica_output_shapes вызовом метода build после инициализации класса tpu embedding.
3. Автоматически определяется из форм входного признака. Приоритет этих мест — тот же, что и порядок.

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_fetch_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=(tpu_features, ))

training_step()

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

enqueue

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

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

Добавляет в очередь тензоры id для поиска встраивания.

Эта функция добавляет в очередь структуру признаков для поиска в таблицах встраивания. Мы ожидаем, что формы входных данных каждого из тензоров в признаках соответствуют формам выходных данных, установленным через FeatureConfig или метод build (если таковые имеются). Формы выходных данных будут автоматически определяться на основе форм входных данных с max_sequence_length или установленным значением output shape в FeatureConfig. Обратите внимание, что формы выходных данных основаны на размере пакета на реплику. Если ваш входной набор данных сгруппирован по глобальному размеру пакета и вы используете experimental_distribute_dataset tf.distribute.TPUStrategy или если вы используете distribute_datasets_from_function и группируете по размеру пакета на ядро, вычисленному контекстом, переданным вашей функции ввода, формы выходных данных должны соответствовать автоматически.

Автоматически определенные формы выходных данных:

1. Для плотных тензоров, если ранг 2 или выше, убедитесь, что тензор имеет последнюю размерность 1. Форма выходных данных будет формой входных данных без последней размерности.
2. Для разреженных тензоров убедитесь, что тензор имеет ранг 2 и выше. а. Если конфигурация признаков имеет max_sequence_length, равное 0, или установлена форма output shape (значение max_sequence_length будет проигнорировано), форма выходных данных будет формой входных данных без последней размерности. б. В противном случае, если тензор имеет ранг 2, форма выходных данных будет формой входных данных с последней размерностью, установленной как max_sequence_length. Если тензор имеет ранг выше 2, форма выходных данных будет формой входных данных без последней размерности, а последняя размерность формы выходных данных будет установлена как max_sequence_length.
3. Для разрозненных тензоров убедитесь, что тензор имеет ранг 2. а. Если конфигурация признаков имеет max_sequence_length, равное 0, или установлена форма output shape (значение max_sequence_length будет проигнорировано), форма выходных данных будет формой входных данных без последней размерности. б. В противном случае форма выходных данных будет формой входных данных без последней размерности, а последняя размерность формы выходных данных будет установлена как max_sequence_length.

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_fetch_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=(tpu_features,))

training_step()

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

Для более тонкого управления в приведенном выше примере строка

embedding.enqueue(embedding_features, training=True)

может быть заменена на

per_core_embedding_features = self.strategy.experimental_local_results(
    embedding_features)

def per_core_enqueue(ctx):
  core_id = ctx.replica_id_in_sync_group
  device = strategy.extended.worker_devices[core_id]
  embedding.enqueue(per_core_embedding_features[core_id],
                    device=device)

strategy.experimental_distribute_values_from_function(
    per_core_queue_inputs)