def test_compose_with_dataset_wires_correctly(self):
with tf.Graph().as_default() as dataset_graph:
d1 = tf.data.Dataset.range(5)
v1 = tf.data.experimental.to_variant(d1)
ds_out_name = v1.name
variant_type = v1.dtype
with tf.Graph().as_default() as reduce_graph:
variant = tf.compat.v1.placeholder(variant_type)
structure = tf.data.experimental.TensorStructure(tf.int64, shape=[])
ds1 = tf.data.experimental.from_variant(variant, structure=structure)
out = ds1.reduce(tf.constant(0, dtype=tf.int64), lambda x, y: x + y)
ds_in_name = variant.name
reduce_out_name = out.name
with dataset_graph.as_default():
init_op_name_1 = tf.initializers.global_variables().name
with reduce_graph.as_default():
init_op_name_2 = tf.initializers.global_variables().name
dataset_graph_spec = graph_merge.GraphSpec(dataset_graph.as_graph_def(),
init_op_name_1, [],
[ds_out_name])
reduce_graph_spec = graph_merge.GraphSpec(reduce_graph.as_graph_def(),
init_op_name_2, [ds_in_name],
[reduce_out_name])
arg_list = [reduce_graph_spec, dataset_graph_spec]
composed_graph, _, _, out_name_map = graph_merge.compose_graph_specs(
arg_list)
with composed_graph.as_default():
with tf.compat.v1.Session() as sess:
ten = sess.run(out_name_map[reduce_out_name])
self.assertEqual(ten, 10)
def test_compose_three_add_one_graphs_adds_three(self):
graph1, input_name_1, output_name_1 = _make_add_one_graph()
graph2, input_name_2, output_name_2 = _make_add_one_graph()
graph3, input_name_3, output_name_3 = _make_add_one_graph()
with graph1.as_default():
init_op_name_1 = tf.compat.v1.global_variables_initializer().name
with graph2.as_default():
init_op_name_2 = tf.compat.v1.global_variables_initializer().name
with graph3.as_default():
init_op_name_3 = tf.compat.v1.global_variables_initializer().name
graph_spec_1 = graph_spec.GraphSpec(graph1.as_graph_def(),
init_op_name_1, [input_name_1],
[output_name_1])
graph_spec_2 = graph_spec.GraphSpec(graph2.as_graph_def(),
init_op_name_2, [input_name_2],
[output_name_2])
graph_spec_3 = graph_spec.GraphSpec(graph3.as_graph_def(),
init_op_name_3, [input_name_3],
[output_name_3])
arg_list = [graph_spec_1, graph_spec_2, graph_spec_3]
composed_graph, init_op_name, in_name_map, out_name_map = graph_merge.compose_graph_specs(
arg_list)
with composed_graph.as_default():
with tf.compat.v1.Session() as sess:
sess.run(init_op_name)
outputs = sess.run(out_name_map[output_name_3],
feed_dict={
in_name_map[input_name_1]: 0.0,
})
self.assertAllClose(outputs, np.array(3.))
def test_compose_no_input_graphs_raises(self):
graph1 = tf.Graph()
with graph1.as_default():
out1 = tf.constant(1.0)
init_op_name_1 = tf.compat.v1.global_variables_initializer().name
graph2 = tf.Graph()
with graph2.as_default():
out2 = tf.constant(2.0)
init_op_name_2 = tf.compat.v1.global_variables_initializer().name
graph_spec_1 = graph_spec.GraphSpec(graph1.as_graph_def(),
init_op_name_1, [], [out1.name])
graph_spec_2 = graph_spec.GraphSpec(graph2.as_graph_def(),
init_op_name_2, [], [out2.name])
arg_list = [graph_spec_1, graph_spec_2]
with self.assertRaisesRegex(ValueError, 'mismatch'):
graph_merge.compose_graph_specs(arg_list)
def test_nested_composition_three_add_one_graphs_adds_three(self):
graph1, input_name_1, output_name_1 = _make_add_variable_number_graph()
graph2, input_name_2, output_name_2 = _make_add_variable_number_graph()
graph3, input_name_3, output_name_3 = _make_add_variable_number_graph()
with graph1.as_default():
init_op_name_1 = tf.compat.v1.global_variables_initializer().name
with graph2.as_default():
init_op_name_2 = tf.compat.v1.global_variables_initializer().name
with graph3.as_default():
init_op_name_3 = tf.compat.v1.global_variables_initializer().name
graph_spec_1 = graph_spec.GraphSpec(graph1.as_graph_def(),
init_op_name_1, [input_name_1],
[output_name_1])
graph_spec_2 = graph_spec.GraphSpec(graph2.as_graph_def(),
init_op_name_2, [input_name_2],
[output_name_2])
arg_list = [graph_spec_1, graph_spec_2]
(partial_merge_graph, partial_merge_init_op_name,
partial_merge_in_name_map, partial_merge_out_name_map
) = graph_merge.compose_graph_specs(arg_list)
partial_graph_spec = graph_spec.GraphSpec(
partial_merge_graph.as_graph_def(), partial_merge_init_op_name,
partial_merge_in_name_map.values(),
partial_merge_out_name_map.values())
graph_spec_3 = graph_spec.GraphSpec(graph3.as_graph_def(),
init_op_name_3, [input_name_3],
[output_name_3])
composed_graph, init_op_name, in_name_map, out_name_map = graph_merge.compose_graph_specs(
[graph_spec_3, partial_graph_spec])
with tf.compat.v1.Session(graph=composed_graph) as sess:
sess.run(init_op_name)
outputs = sess.run(
out_name_map[output_name_3],
feed_dict={
in_name_map[partial_merge_in_name_map[input_name_1]]: 0.0,
})
self.assertAllClose(outputs, np.array(3.))
def test_composition_happens_in_mathematical_composition_order(self):
graph1, input_name_1, output_name_1 = _make_add_one_graph()
def _make_cast_to_int_graph():
with tf.Graph().as_default() as graph:
input_val = tf.compat.v1.placeholder(tf.float32, name='input')
out = tf.cast(input_val, tf.int32)
return graph, input_val.name, out.name
graph2, input_name_2, output_name_2 = _make_cast_to_int_graph()
with graph1.as_default():
init_op_name_1 = tf.compat.v1.global_variables_initializer().name
with graph2.as_default():
init_op_name_2 = tf.compat.v1.global_variables_initializer().name
graph_spec_1 = graph_spec.GraphSpec(graph1.as_graph_def(),
init_op_name_1, [input_name_1],
[output_name_1])
graph_spec_2 = graph_spec.GraphSpec(graph2.as_graph_def(),
init_op_name_2, [input_name_2],
[output_name_2])
arg_list = [graph_spec_2, graph_spec_1]
composed_graph, _, in_name_map, out_name_map = graph_merge.compose_graph_specs(
arg_list)
with composed_graph.as_default():
with tf.compat.v1.Session() as sess:
outputs = sess.run(out_name_map[output_name_2],
feed_dict={
in_name_map[input_name_1]: 0.0,
})
self.assertEqual(outputs, 1)
with self.assertRaises(ValueError):
graph_merge.compose_graph_specs(list(reversed(arg_list)))
def test_compose_two_add_variable_number_graphs_executes_correctly(self):
graph1, input_name_1, output_name_1 = _make_add_variable_number_graph()
graph2, input_name_2, output_name_2 = _make_add_variable_number_graph()
with graph1.as_default():
init_op_name_1 = tf.compat.v1.global_variables_initializer().name
with graph2.as_default():
init_op_name_2 = tf.compat.v1.global_variables_initializer().name
graph_spec_1 = graph_spec.GraphSpec(graph1.as_graph_def(),
init_op_name_1, [input_name_1],
[output_name_1])
graph_spec_2 = graph_spec.GraphSpec(graph2.as_graph_def(),
init_op_name_2, [input_name_2],
[output_name_2])
arg_list = [graph_spec_1, graph_spec_2]
composed_graph, init_op_name, in_name_map, out_name_map = graph_merge.compose_graph_specs(
arg_list)
with tf.compat.v1.Session(graph=composed_graph) as sess:
sess.run(init_op_name)
output_one = sess.run(out_name_map[output_name_2],
feed_dict={
in_name_map[input_name_1]: 0.0,
})
sess.run(init_op_name) # TFF is functional, reset session state.
output_two = sess.run(out_name_map[output_name_2],
feed_dict={
in_name_map[input_name_1]: 0.0,
})
sess.run(init_op_name) # TFF is functional, reset session state.
output_three = sess.run(out_name_map[output_name_2],
feed_dict={
in_name_map[input_name_1]: 0.0,
})
self.assertAllClose(output_one, np.array(2.))
self.assertAllClose(output_two, np.array(2.))
self.assertAllClose(output_three, np.array(2.))
def test_raises_on_list_of_ints(self):
with self.assertRaises(TypeError):
graph_merge.compose_graph_specs([0, 1])
def test_raises_on_graph_spec_set(self):
graph1, input_name_1, output_name_1 = _make_add_one_graph()
graph_spec_1 = graph_spec.GraphSpec(graph1.as_graph_def(), '',
[input_name_1], [output_name_1])
with self.assertRaises(TypeError):
graph_merge.compose_graph_specs(set(graph_spec_1))
def test_raises_on_none(self):
with self.assertRaises(TypeError):
graph_merge.compose_graph_specs(None)
