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。非经特殊声明，原始代码版权归原作者所有，本译文未经允许或授权，请勿转载或复制。