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 等函数中使用表来执行嵌入查找并传递给您的模型。