Running `defun` in graph-mode

问题

I see code like this in TF:

from tensorflow.python.eager import function

...

class _PerDeviceGenerator(dataset_ops.DatasetV2):
  """A `dummy` generator dataset."""

  def __init__(self, shard_num, multi_device_iterator_resource, incarnation_id,
               source_device, element_spec):

    ...

    # TODO(b/124254153): Enable autograph once the overhead is low enough.
    @function.defun(autograph=False)  # Pure graph code.
    def _remote_init_func():
      return functional_ops.remote_call(
          target=source_device,
          args=init_func_concrete.captured_inputs,
          Tout=[dtypes.string],
          f=init_func_concrete)

    self._init_func = _remote_init_func._get_concrete_function_internal()  # pylint: disable=protected-access

    ...

    variant_tensor = gen_dataset_ops.generator_dataset(
        self._init_captured_args,
        self._next_captured_args,
        self._finalize_captured_args,
        init_func=self._init_func,
        next_func=self._next_func,
        finalize_func=self._finalize_func,
        **self._flat_structure)
    super(_PerDeviceGenerator, self).__init__(variant_tensor)

(This code snippet is from TF 1.15.0.)

I'm trying to understand the code.

More specifically, I wonder about defun here. I thought defun is for eager mode.

But here, this code seems to be used for both eager mode and graph mode. Or is that wrong, and this works only on eager mode? (But below, there is MultiDeviceIterator, which has checks like if context.executing_eagerly() and later uses _PerDeviceGenerator for both eager and graph mode. Or is that broken as well for graph mode? Why the check executing_eagerly then?)

What does defun do in graph mode?

That _get_concrete_function_internal is some internal API?

回答1:

Both defun and _get_concrete_function_internal are internal APIs. You should prefer using tf.function whenever possible (or def_function.function when working on internal code). defun is an old API that largely duplicates function, and will likely be removed in the future.

That said, defun is not just for eager mode (neither is function). They both create a TF function, and in graph mode there are "function call" ops that let you call these functions. In eager mode, it just lets you call a TF graph. But in graph mode they let you reduce the size of the graph, in the same way you reduce the size of normal code by factoring common code into functions.

来源：https://stackoverflow.com/questions/61964090/running-defun-in-graph-mode

标签

tensorflow