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

TPUEmbedding 中級 API。

用法

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))

創建要傳遞給入隊操作的分布式數據集時，必須指定特殊的輸入選項：

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

不同的特征輸入可以有不同的形狀。對於密集和稀疏張量，支持 2 級及以上。對於不規則張量，雖然隻支持 rank 2，但您可以將輸出形狀指定為 rank 2 及以上。 FeatureConfig 中指定的輸出形狀具有第一優先級。在構建方法中傳遞的輸入形狀具有第二優先級，從輸入特征自動檢測到的輸入形狀具有最低優先級。後兩者將通過省略最後一個維度轉換為輸出形狀。如果較低優先級的輸出形狀與前一個不匹配。將引發 ValueError。隻有當前者具有未定義的輸出形狀時，後者才能覆蓋。

注意：傳遞給層的所有批次都可以具有不同的輸入形狀。但是這些輸入形狀需要與 FeatureConfig 或 build 方法設置的輸出形狀匹配，不規則張量除外。隻要元素的總數匹配，隻允許輸出形狀設置為更高維度的 2D 不規則張量。所有後續調用必須具有相同的輸入形狀。如果 enqueue 方法無法自動確定輸入形狀，則必須使用輸入形狀調用 build 方法或在FeatureConfig 中提供輸出形狀來初始化圖層。

要在 TPU 上使用此 API，您應該使用自定義訓練循環。以下是訓練和評估步驟的示例：

@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 的調用在使用embedding.apply_gradients 時將training 設置為True，並在函數中不存在embedding.apply_gradients 時設置為False。如果您不遵循此模式，您可能會引發錯誤或 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 等函數中使用表來執行嵌入查找並傳遞給您的模型。