Python tf.xla.experimental.jit_scope用法及代碼示例

在範圍內啟用或禁用運算符的 JIT 編譯。

用法

@contextlib.contextmanager
tf.xla.experimental.jit_scope(
    compile_ops=True, separate_compiled_gradients=False
)

參數

compile_ops 是否啟用或禁用範圍內的編譯。 Python bool 或接受參數 node_def 並返回 python bool 的可調用對象。
separate_compiled_gradients 如果為真，則將每個梯度子圖放入單獨的編譯範圍。這可以細粒度地控製圖形的哪些部分將被編譯為一個單元。單獨編譯梯度可能會為某些圖帶來更好的性能。範圍是根據前向計算的範圍以及梯度的名稱來命名的。因此，梯度將在與前向計算和其他梯度分開的範圍內編譯。

拋出

RuntimeError 如果在啟用即刻執行時調用。

生成(Yield)

當前範圍，啟用或禁用編譯。

注意：這是一個實驗性函數。

編譯是一個提示，僅在盡力而為的基礎上提供支持。

示例用法：

with tf.xla.experimental.jit_scope():
  c = tf.matmul(a, b)  # compiled
with tf.xla.experimental.jit_scope(compile_ops=False):
  d = tf.matmul(a, c)  # not compiled
with tf.xla.experimental.jit_scope(
    compile_ops=lambda node_def:'matmul' in node_def.op.lower()):
  e = tf.matmul(a, b) + d  # matmul is compiled, the addition is not.

separate_compiled_gradients 的示例：

# In the example below, the computations for f, g and h will all be compiled
  # in separate scopes.
  with tf.xla.experimental.jit_scope(
      separate_compiled_gradients=True):
    f = tf.matmul(a, b)
  g = tf.gradients([f], [a, b], name='mygrads1')
  h = tf.gradients([f], [a, b], name='mygrads2')

不在範圍內的操作可能會與 compile_ops=True 範圍內的操作進行集群和編譯，而 compile_ops=False 範圍內的操作將永遠不會被編譯。

例如：

# In the example below, x and loss may be clustered and compiled together,
# while y will not be compiled.
with tf.xla.experimental.jit_scope():
  x = tf.matmul(a, b)
with tf.xla.experimental.jit_scope(compile_ops=False):
  y = tf.matmul(c, d)
loss = x + y

如果您隻想使用 compile_ops=True 編譯範圍內的操作，請考慮添加外部 jit_scope(compile_ops=False) ：

# In the example below, only x will be compiled.
  with tf.xla.experimental.jit_scope(compile_ops=False):
    with tf.xla.experimental.jit_scope():
      x = tf.matmul(a, b)
    y = tf.matmul(c, d)
    loss = x + y

相關用法

注：本文由純淨天空篩選整理自tensorflow.org大神的英文原創作品 tf.xla.experimental.jit_scope。非經特殊聲明，原始代碼版權歸原作者所有，本譯文未經允許或授權，請勿轉載或複製。