tf.xla.experimental.jit_scope
| Просмотреть исходный код на GitHub |
Включить или отключить JIT-компиляцию операторов в области действия.
@contextlib.contextmanager
tf.xla.experimental.jit_scope(
compile_ops=True, separate_compiled_gradients=False
)
Примечание: Это экспериментальная функция.
Компиляция — это подсказка и поддерживается только в лучшем случае.
Пример использования:
with tf.xla.experimental.jit_scope(): c = tf.matmul(a, b) # скомпилировано с tf.xla.experimental.jit_scope(compile_ops=False): d = tf.matmul(a, c) # не скомпилировано с tf.xla.experimental.jit_scope( compile_ops=lambda node_def: 'matmul' in node_def.op.lower()): e = tf.matmul(a, b) + d # matmul скомпилирован, сложение — нет.
Пример separate_compiled_gradients: # В примере ниже вычисления для f, g и h будут все скомпилированы # в отдельных областях. 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 | Включить или отключить компиляцию в области действия. Либо булево значение Python, либо вызываемая функция, которая принимает параметр node_def и возвращает булево значение Python. |
separate_compiled_gradients | Если true, каждый подграф градиента помещается в отдельную область компиляции. Это обеспечивает тонкую настройку, какие части графа будут скомпилированы как единое целое. Отдельная компиляция градиентов может улучшить производительность для некоторых графов. Область действия называется на основе области вычислений вперёд, а также имени градиентов. В результате градиенты будут скомпилированы в области действия, отделенной как от вычислений вперёд, так и от других градиентов. |
| Исключения | |
|---|---|
RuntimeError | если вызывается при включённом режиме eager execution. |
Возвращаемые значения:
Текущая область действия, включая или отключая компиляцию.
© 2020 The TensorFlow Authors. All rights reserved.
Licensed under the Creative Commons Attribution License 3.0.
Code samples licensed under the Apache 2.0 License.
https://www.tensorflow.org/versions/r1.15/api_docs/python/tf/xla/experimental/jit_scope
