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.initializers.global_variables().name
with graph2.as_default():
init_op_name_2 = tf.initializers.global_variables().name
with graph3.as_default():
init_op_name_3 = tf.initializers.global_variables().name
graph_spec_1 = graph_merge.GraphSpec(graph1.as_graph_def(),
init_op_name_1, [input_name_1],
[output_name_1])
graph_spec_2 = graph_merge.GraphSpec(graph2.as_graph_def(),
init_op_name_2, [input_name_2],
[output_name_2])
graph_spec_3 = graph_merge.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_concatenate_inputs_and_outputs_with_dataset_wires_correctly(self):
dataset_graph, _, dataset_out_name = _make_dataset_constructing_graph()
graph_1, _, out_name_1 = _make_manual_reduce_graph(
dataset_graph, dataset_out_name)
graph_2, _, out_name_2 = _make_manual_reduce_graph(
dataset_graph, dataset_out_name)
with graph_1.as_default():
init_op_name_1 = tf.initializers.global_variables().name
with graph_2.as_default():
init_op_name_2 = tf.initializers.global_variables().name
graph_spec_1 = graph_merge.GraphSpec(graph_1.as_graph_def(),
init_op_name_1, [], [out_name_1])
graph_spec_2 = graph_merge.GraphSpec(graph_2.as_graph_def(),
init_op_name_2, [], [out_name_2])
arg_list = [graph_spec_1, graph_spec_2]
merged_graph, init_op_name, _, out_name_maps = graph_merge.concatenate_inputs_and_outputs(
arg_list)
with merged_graph.as_default():
with tf.compat.v1.Session() as sess:
sess.run(init_op_name)
tens = sess.run([
out_name_maps[0][out_name_1], out_name_maps[1][out_name_2]
])
self.assertEqual(tens, [10, 10])
def test_concatenate_inputs_and_outputs_no_arg_graphs(self):
graph1 = tf.Graph()
with graph1.as_default():
out1 = tf.constant(1.0)
init_op_name_1 = tf.initializers.global_variables().name
graph2 = tf.Graph()
with graph2.as_default():
out2 = tf.constant(2.0)
init_op_name_2 = tf.initializers.global_variables().name
graph_spec_1 = graph_merge.GraphSpec(graph1.as_graph_def(),
init_op_name_1, [], [out1.name])
graph_spec_2 = graph_merge.GraphSpec(graph2.as_graph_def(),
init_op_name_2, [], [out2.name])
arg_list = [graph_spec_1, graph_spec_2]
merged_graph, init_op_name, _, out_name_maps = graph_merge.concatenate_inputs_and_outputs(
arg_list)
with merged_graph.as_default():
with tf.compat.v1.Session() as sess:
sess.run(init_op_name)
outputs = sess.run(
[out_name_maps[0][out1.name], out_name_maps[1][out2.name]])
self.assertAllClose(outputs, np.array([1., 2.]))
def test_concatenate_inputs_and_outputs_no_init_op_graphs(self):
graph1, input_name_1, output_name_1 = _make_add_one_graph()
graph2, input_name_2, output_name_2 = _make_add_one_graph()
graph_spec_1 = graph_merge.GraphSpec(graph1.as_graph_def(), None,
[input_name_1], [output_name_1])
graph_spec_2 = graph_merge.GraphSpec(graph2.as_graph_def(), None,
[input_name_2], [output_name_2])
arg_list = [graph_spec_1, graph_spec_2]
merged_graph, init_op_name, in_name_maps, out_name_maps = graph_merge.concatenate_inputs_and_outputs(
arg_list)
with merged_graph.as_default():
with tf.compat.v1.Session() as sess:
sess.run(init_op_name)
outputs = sess.run(
[
out_name_maps[0][output_name_1],
out_name_maps[1][output_name_2]
],
feed_dict={
in_name_maps[0][input_name_1]: 1.0,
in_name_maps[1][input_name_2]: 2.0
})
self.assertAllClose(outputs, np.array([2., 3.]))
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_concatenate_inputs_and_outputs_three_add_one_graphs(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.initializers.global_variables().name
with graph2.as_default():
init_op_name_2 = tf.initializers.global_variables().name
with graph3.as_default():
init_op_name_3 = tf.initializers.global_variables().name
graph_spec_1 = graph_merge.GraphSpec(graph1.as_graph_def(), init_op_name_1,
[input_name_1], [output_name_1])
graph_spec_2 = graph_merge.GraphSpec(graph2.as_graph_def(), init_op_name_2,
[input_name_2], [output_name_2])
graph_spec_3 = graph_merge.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]
merged_graph, init_op_name, in_name_maps, out_name_maps = graph_merge.concatenate_inputs_and_outputs(
arg_list)
with merged_graph.as_default():
with tf.compat.v1.Session() as sess:
sess.run(init_op_name)
outputs = sess.run(
[
out_name_maps[0][output_name_1],
out_name_maps[1][output_name_2], out_name_maps[2][output_name_3]
],
feed_dict={
in_name_maps[0][input_name_1]: 1.0,
in_name_maps[1][input_name_2]: 2.0,
in_name_maps[2][input_name_3]: 3.0
})
self.assertAllClose(outputs, np.array([2., 3., 4.]))
def test_concatenate_inputs_and_outputs_two_add_variable_number_graphs(
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_merge.GraphSpec(graph1.as_graph_def(),
init_op_name_1, [input_name_1],
[output_name_1])
graph_spec_2 = graph_merge.GraphSpec(graph2.as_graph_def(),
init_op_name_2, [input_name_2],
[output_name_2])
arg_list = [graph_spec_1, graph_spec_2]
merged_graph, init_op_name, in_name_maps, out_name_maps = graph_merge.concatenate_inputs_and_outputs(
arg_list)
with merged_graph.as_default():
with tf.compat.v1.Session() as sess:
sess.run(init_op_name)
outputs_1 = sess.run(
[
out_name_maps[0][output_name_1],
out_name_maps[1][output_name_2]
],
feed_dict={
in_name_maps[0][input_name_1]: 1.0,
in_name_maps[1][input_name_2]: 2.0
})
outputs_2 = sess.run(
[
out_name_maps[0][output_name_1],
out_name_maps[1][output_name_2]
],
feed_dict={
in_name_maps[0][input_name_1]: 1.0,
in_name_maps[1][input_name_2]: 2.0
})
outputs_3 = sess.run(
[
out_name_maps[0][output_name_1],
out_name_maps[1][output_name_2]
],
feed_dict={
in_name_maps[0][input_name_1]: 1.0,
in_name_maps[1][input_name_2]: 2.0
})
self.assertAllClose(outputs_1, [2., 3.])
self.assertAllClose(outputs_2, [3., 4.])
self.assertAllClose(outputs_3, [4., 5.])
def _unpack_proto_into_graph_spec(tf_block_proto):
"""Packs a TF proto into a `graph_merge.GraphSpec`.
Args:
tf_block_proto: Instance of `computation_pb2.Computation` with `tensorflow`
`computation` attribute.
Returns:
Instance of `graph_merge.GraphSpec` containing Python representations of
the information present in `tf_block_proto`.
"""
graph = serialization_utils.unpack_graph_def(
tf_block_proto.tensorflow.graph_def)
graph_init_op_name = tf_block_proto.tensorflow.initialize_op
if not graph_init_op_name:
graph_init_op_name = None
graph_parameter_binding = tf_block_proto.tensorflow.parameter
graph_result_binding = tf_block_proto.tensorflow.result
if graph_parameter_binding.WhichOneof('binding') is not None:
graph_parameter_list = graph_utils.extract_tensor_names_from_binding(
graph_parameter_binding)
else:
graph_parameter_list = []
graph_result_list = graph_utils.extract_tensor_names_from_binding(
graph_result_binding)
return graph_merge.GraphSpec(graph, graph_init_op_name,
graph_parameter_list, graph_result_list)
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.initializers.global_variables().name
graph2 = tf.Graph()
with graph2.as_default():
out2 = tf.constant(2.0)
init_op_name_2 = tf.initializers.global_variables().name
graph_spec_1 = graph_merge.GraphSpec(graph1.as_graph_def(),
init_op_name_1, [], [out1.name])
graph_spec_2 = graph_merge.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_graph_spec_constructs_dummy_data(self):
graph_def = _make_add_one_graph()[0].as_graph_def()
init_op = 'init'
in_names = ['in']
out_names = ['out']
x = graph_merge.GraphSpec(graph_def, init_op, in_names, out_names)
self.assertIs(x.graph_def, graph_def)
self.assertIs(x.init_op, init_op)
self.assertIs(x.in_names, in_names)
self.assertIs(x.out_names, out_names)
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.initializers.global_variables().name
with graph2.as_default():
init_op_name_2 = tf.initializers.global_variables().name
graph_spec_1 = graph_merge.GraphSpec(graph1.as_graph_def(),
init_op_name_1, [input_name_1],
[output_name_1])
graph_spec_2 = graph_merge.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.initializers.global_variables().name
with graph2.as_default():
init_op_name_2 = tf.initializers.global_variables().name
graph_spec_1 = graph_merge.GraphSpec(graph1.as_graph_def(), init_op_name_1,
[input_name_1], [output_name_1])
graph_spec_2 = graph_merge.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 composed_graph.as_default():
with tf.compat.v1.Session() 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,
})
output_two = sess.run(
out_name_map[output_name_2],
feed_dict={
in_name_map[input_name_1]: 0.0,
})
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(4.))
self.assertAllClose(output_three, np.array(6.))
def test_graph_spec_fails_bad_init_op(self):
graph_def = _make_add_one_graph()[0].as_graph_def()
with self.assertRaises(TypeError):
graph_merge.GraphSpec(graph_def, 1, ['test'], ['test'])
def test_graph_spec_fails_no_graph_def(self):
with self.assertRaises(TypeError):
graph_merge.GraphSpec(None, 'test', ['test'], ['test'])
def test_graph_spec_succeeds_empty_init_op(self):
graph_def = _make_add_one_graph()[0].as_graph_def()
graph_merge.GraphSpec(graph_def, '', ['test'], ['test'])
def test_graph_spec_fails_no_out_names(self):
graph_def = _make_add_one_graph()[0].as_graph_def()
with self.assertRaises(TypeError):
graph_merge.GraphSpec(graph_def, 'test', ['test'], None)
def test_raises_on_graph_spec_set(self):
graph1, input_name_1, output_name_1 = _make_add_one_graph()
graph_spec_1 = graph_merge.GraphSpec(graph1.as_graph_def(), '',
[input_name_1], [output_name_1])
with self.assertRaises(TypeError):
graph_merge.compose_graph_specs(set(graph_spec_1))
