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 对象的实例。