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)