I found the response here.
Actually, you can enable graph export in v2. You'll need to call
tf.summary.trace_on()
before the code you want to trace the graph for
(e.g. L224 if you just want the train step), and then call
tf.summary.trace_off()
after the code completes. Since you only need
one trace of the graph, I would recommend wrapping these calls with if
global_step_val == 0:
so that you don't produce traces every step.
Actually, to create the graph it is necessary to do the trace just once and makes no sense of doing it at each epoch. The solution is just to check before calling the trace just once like:
for epoch in range(epochs):
if epoch == 0:
tf.summary.trace_on(graph=True, profiler=True)
z = train_step(x, y)
if epoch == 0:
with writer.as_default():
tf.summary.trace_export(name="train_graph", step=0, profiler_outdir=logdir)
I personally like more this decorator idea:
def run_once(f):
def wrapper(*args, **kwargs):
if not wrapper.has_run:
wrapper.has_run = True
return f(*args, **kwargs)
wrapper.has_run = False
return wrapper
@run_once
def _start_graph_tensorflow(self):
tf.summary.trace_on(graph=True, profiler=True) # https://www.tensorflow.org/tensorboard/graphs
@run_once
def _end_graph_tensorflow(self):
with self.graph_writer.as_default():
tf.summary.trace_export(name="graph", step=0, profiler_outdir=self.graph_writer_logdir)
for epoch in range(epochs):
_start_graph_tensorflow()
z = train_step(x, y)
_end_graph_tensorflow()