python - 如何在 tensorboard 中显示 Tensorflow 2.0 中的 tf.data.Dataset.map 子图？

Question

根据documentation , tf.data.Datasets在图形模式下工作(在 Eager 和图形模式下):

Note that irrespective of the context in which map_func is defined (eager vs. graph), tf.data traces the function and executes it as a graph

在 Tensorflow 1.X 中，我们可以轻松地在 Tensorboard 中绘制此图:处理函数绘制在子图中。

例如，

def _parse_function(x):
    return x * 2

x = tf.constant([0 , 1])
dataset = tf.data.Dataset.from_tensor_slices(x)
dataset = dataset.map(_parse_function)

在 Tensorboard 中，会出现一个子图，对应于 _parse_function:


然而，在 Tensorflow 2.0 中，这不会在 Tensorboard 图中产生任何可见元素。
根据 Tensorboard，以下代码不会生成任何图形:

def _parse_function(x):
    return x * 2

logdir = 'logs'
writer = tf.summary.create_file_writer(logdir)

tf.summary.trace_on(graph=True, profiler=True)
x = tf.constant([0 , 1])
dataset = tf.data.Dataset.from_tensor_slices(x)
dataset = dataset.map(_parse_function)

with writer.as_default():
  tf.summary.trace_export(
      name="trace",
      step=0,
      profiler_outdir=logdir)

因此，由于在调用 map 时创建了图形。 , 有没有办法访问/可视化这个图？

Answer 1

在函数内部创建操作并用 tf.function 装饰它

import tensorflow as tf
def _parse_function(x):
    return x * 2

@tf.function
def foo():
    x = tf.constant([0 , 1])
    dataset = tf.data.Dataset.from_tensor_slices(x)
    dataset = dataset.map(_parse_function)

logdir = 'logs'
writer = tf.summary.create_file_writer(logdir)

tf.summary.trace_on(graph=True, profiler=True)
foo()
with writer.as_default():
    tf.summary.trace_export(
      name="trace",
      step=0,
      profiler_outdir=logdir)