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. Ниже приведены примеры использования этих методов для обучения и оценки вашей модели. На процессоре доступна только свойство embedding_tables, предоставляющая доступ к таблицам встраивания, что позволяет использовать их для оценки/предсказания модели на процессоре.

Сначала рассмотрим режим 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, за исключением разрозненных тензоров. Разрешены только разрозненные тензоры ранга 2 с формой вывода, заданной для более высоких измерений, при условии, что общее количество элементов совпадает. Все последующие вызовы должны иметь одинаковые формы ввода. В случае, если формы ввода не могут быть автоматически определены методом 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 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=(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(...)

На процессоре доступна только свойство 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. Конфигурация признаков, предоставленная в функции init.
2. Per_replica_output_shapes путём непосредственного вызова метода build после инициализации класса tpu встраивания.
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 или установки формы вывода в FeatureConfig. Обратите внимание, что форма вывода основана на размере пакета на реплику. Если ваш входной набор данных сгруппирован по глобальному размеру пакета и вы используете tf.distribute.TPUStrategy's experimental_distribute_dataset, или если вы используете distribute_datasets_from_function и группируете по размеру пакета на ядро, вычисленный контекстом, переданным вашей функции ввода, формы вывода должны автоматически соответствовать.

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

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