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 и выше. Для ragged-тензоров поддерживается только ранг 2, однако выходную форму можно задать с рангом 2 и выше. Приоритет при выборе выходной формы в первую очередь имеет значение, заданное в FeatureConfig. Второй по приоритету является входная форма, переданная в метод build, а наименьший приоритет имеют входные формы, автоматически определенные по входному признаку. Два последних варианта преобразуются в выходные формы путем удаления последнего измерения. Если выходные формы с меньшим приоритетом не соответствуют формам с более высоким приоритетом, будет вызвано исключение ValueError. Переопределение значений с более высоким приоритетом допускается только в том случае, если для них не определены выходные формы.

Примечание: Все пакеты, передаваемые слою, могут иметь разные входные формы. Однако эти формы должны соответствовать выходным формам, заданным в FeatureConfig или методе build, за исключением ragged-тензоров. Допускаются только двумерные ragged-тензоры с выходной формой более высокого ранга при условии, что совпадает общее количество элементов. Во всех последующих вызовах должны использоваться те же входные формы. Если метод 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. Выходную форму можно задать в трех местах.

1. В конфигурации FeatureConfig, переданной функции 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
)

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

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

Выходные формы определяются автоматически следующим образом:

1. Для плотного тензора ранга 2 или выше необходимо убедиться, что размер последнего измерения равен 1. Выходная форма будет совпадать с входной формой без последнего измерения.
2. Для разреженного тензора необходимо убедиться, что его ранг равен 2 или выше. a. Если max_sequence_length в конфигурации признака равен 0 или задана выходная форма (параметр max_sequence_length будет проигнорирован), выходная форма будет совпадать с входной формой без последнего измерения. b. В противном случае, если ранг тензора равен 2, выходная форма будет совпадать с входной формой, за исключением того, что последнее измерение будет равно max_sequence_length. Если ранг тензора больше 2, выходная форма будет совпадать с входной формой без последнего измерения, а последнее измерение выходной формы будет равно max_sequence_length.
3. Для ragged-тензора необходимо убедиться, что его ранг равен 2. a. Если max_sequence_length в конфигурации признака равен 0 или задана выходная форма (параметр max_sequence_length будет проигнорирован), выходная форма будет совпадать с входной формой без последнего измерения. b. В противном случае выходная форма будет совпадать с входной формой без последнего измерения, а последнее измерение выходной формы будет равно 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()

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

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

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)