Skip to content

Instantly share code, notes, and snippets.

@tomhennigan
Created February 5, 2019 16:02
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 tomhennigan/a183a6570d2a6fb9bce08758ee2d78b6 to your computer and use it in GitHub Desktop.
Save tomhennigan/a183a6570d2a6fb9bce08758ee2d78b6 to your computer and use it in GitHub Desktop.
import tensorflow as tf
tf.enable_v2_behavior()
l = snt.Linear(64)
@tf.function(autograph=False)
def f(x):
y1 = l(x)
y2 = l(x)
return y1, y2
g = f.get_concrete_function(tf.TensorSpec([1, 64], tf.float32)).graph
for node in g.as_graph_def().node:
print node.name
print
g = f.get_concrete_function(tf.TensorSpec([2, 64], tf.float32)).graph
for node in g.as_graph_def().node:
print node.name
x
linear/MatMul/ReadVariableOp/resource
linear/MatMul/ReadVariableOp
linear/MatMul
linear/Add/ReadVariableOp/resource
linear/Add/ReadVariableOp
linear/Add
linear/MatMul_1/ReadVariableOp
linear/MatMul_1
linear/Add_1/ReadVariableOp
linear/Add_1
Identity
Identity_1
x
linear/MatMul/ReadVariableOp/resource
linear/MatMul/ReadVariableOp
linear/MatMul
linear/Add/ReadVariableOp/resource
linear/Add/ReadVariableOp
linear/Add
linear/MatMul_1/ReadVariableOp
linear/MatMul_1
linear/Add_1/ReadVariableOp
linear/Add_1
Identity
Identity_1
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment