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

Posted 2023-03-27

【中文标题】如何在 tensorboard 中显示 Tensorflow 2.0 中的 tf.data.Dataset.map 子图？

【英文标题】：How to show in tensorboard the tf.data.Dataset.map subgraph in Tensorflow 2.0?

【发布时间】：2019-12-22 22:53:34

【问题描述】：

根据documentation，tf.data.Datasets 在图形模式下工作（在 Eager 和图形模式下）：

请注意，无论定义 map_func 的上下文如何（eager vs. 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时会创建一个图表，有没有办法访问/可视化这个图表？

【参考方案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)