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))