Tensorflow: How to replace a node in a calculation graph?
TL;DR: If you can define the two computations as Python functions, you should do that. If you can't, there's more advanced functionality in TensorFlow to serialize and import graphs, which allows you to compose graphs from different sources.
One way to do this in TensorFlow is to build the disjoint computations as separate tf.Graph
objects, then convert them to serialized protocol buffers using Graph.as_graph_def()
:
with tf.Graph().as_default() as g_1:
input = tf.placeholder(tf.float32, name="input")
y = f(input)
# NOTE: using identity to get a known name for the output tensor.
output = tf.identity(y, name="output")
gdef_1 = g_1.as_graph_def()
with tf.Graph().as_default() as g_2: # NOTE: g_2 not g_1
input = tf.placeholder(tf.float32, name="input")
z = g(input)
output = tf.identity(y, name="output")
gdef_2 = g_2.as_graph_def()
Then you could compose gdef_1
and gdef_2
into a third graph, using tf.import_graph_def()
:
with tf.Graph().as_default() as g_combined:
x = tf.placeholder(tf.float32, name="")
# Import gdef_1, which performs f(x).
# "input:0" and "output:0" are the names of tensors in gdef_1.
y, = tf.import_graph_def(gdef_1, input_map={"input:0": x},
return_elements=["output:0"])
# Import gdef_2, which performs g(y)
z, = tf.import_graph_def(gdef_2, input_map={"input:0": y},
return_elements=["output:0"]
If you want to combine trained models (for example to reuse parts of a pretrained model in a new model), you can use a Saver
to save a checkpoint of the first model, then restore that model (entirely or partially) into another model.
For example, say you want to reuse model 1's weights w
in model 2, and also convert x
from a placeholder to a variable:
with tf.Graph().as_default() as g1:
x = tf.placeholder('float')
w = tf.Variable(1., name="w")
y = x * w
saver = tf.train.Saver()
with tf.Session(graph=g1) as sess:
w.initializer.run()
# train...
saver.save(sess, "my_model1.ckpt")
with tf.Graph().as_default() as g2:
x = tf.Variable(2., name="v")
w = tf.Variable(0., name="w")
z = x + w
restorer = tf.train.Saver([w]) # only restore w
with tf.Session(graph=g2) as sess:
x.initializer.run() # x now needs to be initialized
restorer.restore(sess, "my_model1.ckpt") # restores w=1
print(z.eval()) # prints 3.
It turns out that tf.train.import_meta_graph
passes all additional arguments to the underlying import_scoped_meta_graph
which has the input_map
argument and utilizes it when it gets to it's own (internal) invocation of import_graph_def
.
It is not documented, and took me waaaay toooo much time to find it, but it works!