Python tf.tpu.experimental.embedding.TPUEmbedding.enqueue用法及代碼示例

用法

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

將 id 張量排入隊列以進行嵌入查找。

此函數將要在嵌入表中查找的特征結構排入隊列。我們期望特征中每個張量的輸入形狀與通過 FeatureConfig 或構建方法(如果有)設置的輸出形狀相匹配。輸出形狀將根據帶有max_sequence_length 的輸入形狀或 FeatureConfig 中的輸出形狀設置自動檢測。請注意，輸出形狀基於每個副本的批量大小。如果您的輸入數據集被批處理為全局批處理大小並且您使用 tf.distribute.TPUStrategy 的 experimental_distribute_dataset 或者如果您使用 distribute_datasets_from_function 並批處理到由傳遞給輸入函數的上下文計算的每個核心批處理大小，則輸出形狀應該自動匹配。

自動檢測到輸出形狀：

1. 對於稠密張量，如果 rank 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()

注意：如上所述使用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)