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

Question

根据文档，tf.data.Dataset以图模式工作（在eager和graph模式下都是如此）：

请注意，无论map_func定义的上下文是eager还是graph，tf.data都会跟踪该函数并将其作为图形执行

在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时会创建图表，是否有一种方式可以访问/可视化此图表？

- Phylliade

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)

- Amal Roy