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

用法

apply_gradients(
    gradients,
    name:Optional[Text] = None
)

参数

gradients 渐变的嵌套结构，其结构与传递给此对象的feature_config 匹配。
name 底层操作的名称。

抛出

RuntimeError 如果在 TPUStrategy 下未创建对象或未构建对象时调用(通过手动调用 build 或调用 enqueue)。
ValueError 如果传入了非tf.Tensor非None梯度，或者传入了不正确形状的tf.Tensor。另外，如果gradients中任何序列的大小与feature_config中的对应序列不匹配.
TypeError 如果 gradients 中任何序列的类型与 feature_config 中的对应序列不匹配。

将梯度更新应用于嵌入表。

如果在嵌套结构的任何位置传递None 的梯度，则对该特征应用具有零梯度的梯度更新。对于像 SGD 或 Adagrad 这样的优化器，这与根本不应用更新是一样的。对于惰性 Adam 和其他具有衰减的稀疏应用优化器，请确保您了解应用零梯度的效果。

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

相关用法

注：本文由纯净天空筛选整理自tensorflow.org大神的英文原创作品 tf.tpu.experimental.embedding.TPUEmbedding.apply_gradients。非经特殊声明，原始代码版权归原作者所有，本译文未经允许或授权，请勿转载或复制。