tf.types.experimental.PolymorphicFunction

Базовый класс для полиморфных графовых функций.

Наследует: Callable

Просмотреть псевдонимы

Основные псевдонимы

tf.types.experimental.GenericFunction

Графовые функции — это вызываемые объекты Python, которые перенаправляют вызовы в граф TensorFlow. Полиморфные графовые функции могут быть основаны на нескольких графах TF и автоматически выбирать подходящую специализацию в зависимости от типа входных данных, с которыми их вызвали. При необходимости они также могут создавать специализации на лету, например посредством трассировки.

См. также tf.function.

Атрибуты
`function_type`	Возвращает FunctionType, описывающий этот вызываемый объект.

Методы

`experimental_get_compiler_ir`

Просмотреть исходный код

experimental_get_compiler_ir(
    *args, **kwargs
)

Возвращает IR компилятора для скомпилированной функции.

Этот API предназначен только для отладки, поскольку обратная совместимость возвращаемого IR и допустимых значений stage не гарантируется.

Аргументы
`*args`	Аргументы компиляции принимают входные данные либо: (1) все входные данные являются TensorSpec, либо (2) все входные данные являются tf.Tensor/переменными Python.
`**kwargs`	Именованные аргументы, используемые для компиляции. Требования те же, что и для аргументов компиляции.

Возвращает

Вызываемый объект функции со следующими именованными аргументами:

stage, на котором следует сериализовать IR компилятора. Допустимые значения:
- hlo: выходные данные HLO после преобразования из TF (https://www.tensorflow.org/xla/operation_semantics).
- hlo_serialized: аналогично stage=hlo, но выходные данные представляют собой сериализованный протокол буфера модуля HLO (объект bytes).
- optimized_hlo: HLO после оптимизаций компилятора.
- optimized_hlo_serialized: аналогично stage=optimized_hlo, но выходные данные представляют собой сериализованный протокол буфера модуля HLO (объект bytes).
- optimized_hlo_dot: оптимизированный HLO в формате DOT, подходящем для Graphviz.
device_name может быть равен None — в этом случае для компиляции используется предпочтительное устройство — или содержать имя устройства. Это может быть полное или частичное имя устройства, например /device:CPU:0.

Например, для

@tf.function(jit_compile=True)
def f(x):
  return x + 1

f.experimental_get_compiler_ir(tf.random.normal([10, 10])(stage='hlo')

результат будет следующим:

HloModule a_inference_f_13__.9

ENTRY %a_inference_f_13__.9 (arg0.1: f32[10,10]) -> f32[10,10] {
  %arg0.1 = f32[10,10]{1,0} parameter(0), parameter_replication={false}
  %reshape.2 = f32[10,10]{1,0} reshape(f32[10,10]{1,0} %arg0.1)
  %constant.3 = f32[] constant(1)
  %broadcast.4 = f32[10,10]{1,0} broadcast(f32[] %constant.3)
  %add.5 = f32[10,10]{1,0} add(f32[10,10]{1,0} %reshape.2,
                               f32[10,10]{1,0} %broadcast.4)
  %reshape.6 = f32[10,10]{1,0} reshape(f32[10,10]{1,0} %add.5)
  %tuple.7 = (f32[10,10]{1,0}) tuple(f32[10,10]{1,0} %reshape.6)
  ROOT %get-tuple-element.8 = f32[10,10]{1,0}
    get-tuple-element((f32[10,10]{1,0}) %tuple.7), index=0
}

Вот еще один пример с входными данными tf.TensorSpec:

y = tf.Variable(tf.zeros([10, 20], dtype=tf.float32))

@tf.function(jit_compile=True)
def f(x):
  return x + y

hlo_str = f.experimental_get_compiler_ir(tf.TensorSpec(shape=(10,
20)))(stage='hlo')

Результат:

HloModule a_inference_f_120__.8,
entry_computation_layout={(f32[10,20]{1,0},f32[10,20]{1,0})->f32[10,20]{1,0} }

ENTRY %a_inference_f_120__.8 (arg0.1: f32[10,20], arg1.2: f32[10,20]) ->
f32[10,20] {
  %arg0.1 = f32[10,20]{1,0} parameter(0), parameter_replication={false},
  metadata={op_name="XLA_Args"}
  %reshape.3 = f32[10,20]{1,0} reshape(f32[10,20]{1,0} %arg0.1)
  %arg1.2 = f32[10,20]{1,0} parameter(1), parameter_replication={false},
  metadata={op_name="XLA_Args"}
  %add.4 = f32[10,20]{1,0} add(f32[10,20]{1,0} %reshape.3, f32[10,20]{1,0}
  %arg1.2), metadata={op_type="AddV2" op_name="add"
  source_file="<ipython-input-16-ea04879c1873>" source_line=4}
  %reshape.5 = f32[10,20]{1,0} reshape(f32[10,20]{1,0} %add.4),
  metadata={op_name="XLA_Retvals"}
  %tuple.6 = (f32[10,20]{1,0}) tuple(f32[10,20]{1,0} %reshape.5),
  metadata={op_name="XLA_Retvals"}
  ROOT %get-tuple-element.7 = f32[10,20]{1,0}
  get-tuple-element((f32[10,20]{1,0}) %tuple.6), index=0,
  metadata={op_name="XLA_Retvals"}
}
</td>
</tr>

</table>

Модуль HLO принимает плоский список входных данных. Чтобы получить порядок этих входных данных в сигнатурах, можно вызвать concrete_fn.structured_input_signature и concrete_fn.captured_inputs:

# Use concrete_fn to get the hlo_module flat_args.
concrete_fn = f.get_concrete_function(tf.TensorSpec(shape=(10, 20)))
flat_args = list(
    tf.nest.flatten(concrete_fn.structured_input_signature)
    ) + concrete_fn.captured_inputs

Вызывает исключения
`ValueError`	(1) Если выбрано недопустимое значение `stage`; (2) если применяется к функции, которая не скомпилирована (`jit_compile=True` не задано); (3) если формы входных данных для tf.TensorSpec не заданы полностью.
`TypeError`	При вызове с входными данными в режиме графа.

`get_concrete_function`

Просмотреть исходный код

@abc.abstractmethod
get_concrete_function(
    *args, **kwargs
) -> tf.types.experimental.ConcreteFunction

Возвращает ConcreteFunction, специализированную для типов входных данных.

Аргументы, заданные с помощью args и kwargs, следуют обычным правилам вызова функций. У возвращаемого объекта ConcreteFunction тот же набор позиционных и именованных аргументов, что и у self, но типы этих аргументов совместимы с типами, заданными с помощью args и kwargs (хотя необязательно совпадают с ними).

@tf.function
def f(x):
  return x
f_concrete = f.get_concrete_function(tf.constant(1.0))
f_concrete = f.get_concrete_function(x=tf.constant(1.0))

В отличие от обычных вызовов, get_concrete_function допускают использование спецификаторов типов вместо объектов TensorFlow. Например, объекты tf.Tensor можно заменить на tf.TensorSpec.

@tf.function
def f(x):
  return x
f_concrete = f.get_concrete_function(tf.TensorSpec([], tf.float64))

Если определение функции допускает только одну специализацию, args и kwargs можно не указывать.

@tf.function(input_signature=[tf.TensorSpec(None, tf.float32)])
def f(x):
  return x
f_concrete = f.get_concrete_function()

Возвращаемый объект ConcreteFunction можно вызывать обычным образом:

f_concrete(tf.constant(1.0))
<tf.Tensor: shape=(), dtype=float32, numpy=1.0>
f_concrete(x=tf.constant(1.0))
<tf.Tensor: shape=(), dtype=float32, numpy=1.0>

Аргументы
`*args`	входные данные для специализации.
`**kwargs`	входные данные для специализации.

Возвращает
`ConcreteFunction`.

`call`

Просмотреть исходный код

__call__(
    *args, **kwargs
)

Выполняет этот вызываемый объект.

Работает как обычный оператор: в eager-режиме выполнение начинается немедленно и возвращает результаты. В режиме графа создаются операторы, возвращающие символические значения TensorFlow (например, tf.Tensor, tf.data.Dataset и т. д.). Например, вызываемые объекты tf.function обычно создают оператор tf.raw_ops.PartitionedCall, но не всегда: конкретные создаваемые операции являются деталью внутренней реализации.

Аргументы
`*args`	позиционный аргумент этого вызова
`**kwargs`	именованные аргументы этого вызова

Возвращает
Результаты выполнения.

Если не указано иное, содержимое этой страницы распространяется по лицензии Creative Commons Attribution 4.0, а примеры кода — по лицензии Apache 2.0. Подробную информацию см. в правилах сайта Google Developers. Java является зарегистрированным товарным знаком Oracle и/или ее аффилированных лиц. Часть содержимого распространяется по лицензии numpy.

Последнее обновление: 2024-04-26 UTC.

© 2022 The TensorFlow Authors. All rights reserved.
Licensed under the Creative Commons Attribution License 4.0.
Code samples licensed under the Apache 2.0 License.
https://www.tensorflow.org/api_docs/python/tf/types/experimental/PolymorphicFunction