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

用法

dequeue(
    name:Optional[Text] = None
)

獲取嵌入結果。

返回 tf.Tensor 對象的嵌套結構，將 feature_config 參數的結構與 TPUEmbedding 類匹配。張量的輸出形狀是 (*output_shape, dim) , dim 是對應的 TableConfig 的維度。對於output_shape，可以設置三個位置。

1. FeatureConfig 在 init 函數中提供。
2. Per_replica_output_shapes通過初始化tpu嵌入類後直接調用build方法。
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 對象的實例。