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。非經特殊聲明，原始代碼版權歸原作者所有，本譯文未經允許或授權，請勿轉載或複製。