Skip to content

Instantly share code, notes, and snippets.

@brandonwillard
Last active July 12, 2019 16:48
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save brandonwillard/cf7df784c69ae4e9718d070e849a62ab to your computer and use it in GitHub Desktop.
Save brandonwillard/cf7df784c69ae4e9718d070e849a62ab to your computer and use it in GitHub Desktop.
Getting a TensorFlow graph in eager-mode
import tensorflow as tf
# from tensorflow.python.framework import ops
# ops.disable_eager_execution()
# Make sure we're in eager-mode
assert tf.executing_eagerly()
@tf.function
def some_function():
"""An example of tf.function."""
A = tf.compat.v1.placeholder(tf.float64, name='A',
shape=tf.TensorShape([None, None]))
x = tf.compat.v1.placeholder(tf.float64, name='x',
shape=tf.TensorShape([None, 1]))
y = tf.compat.v1.placeholder(tf.float64, name='y',
shape=tf.TensorShape([None, 1]))
z = tf.matmul(A, x + y, name='z')
return z
# The graph is available through this function
f_concrete = some_function.get_concrete_function()
# FYI: If `some_function` had arguments...
# f_concrete = some_function.get_concrete_function(tf.TensorSpec([], tf.float64))
# Here's one graph we can use
fgraph = f_concrete.graph
# Everything's there
list(fgraph.outputs[0].op.inputs)
fgraph.outputs[0].op.inputs[0].op
list(fgraph.outputs[0].op.inputs[0].op.inputs)
# There's also this one, which is the "original" V1 graph type.
fgraph.outer_graph
# Looks like we can also use this context manager
from tensorflow.python.eager.context import graph_mode
assert tf.executing_eagerly()
with graph_mode():
A = tf.compat.v1.placeholder(tf.float64, name='A',
shape=tf.TensorShape([None, None]))
x = tf.compat.v1.placeholder(tf.float64, name='x',
shape=tf.TensorShape([None, 1]))
y = tf.compat.v1.placeholder(tf.float64, name='y',
shape=tf.TensorShape([None, 1]))
z = tf.matmul(A, x + y, name='z')
assert tf.executing_eagerly()
z.op
z.op.inputs[0].op
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment