def testAddCellSubgraphSpecHook(self): component = MockComponent() cell = export_pb2.CellSubgraphSpec() cell.input.add( name='feature', tensor='feature_tensor', type=export_pb2.CellSubgraphSpec.Input.TYPE_FEATURE) cell.input.add( name='recurrent', tensor='recurrent_tensor', type=export_pb2.CellSubgraphSpec.Input.TYPE_RECURRENT) cell.output.add(name='layer_0', tensor='layer_0_tensor') cell.output.add(name='logits', tensor='logits_tensor') with self.test_session() as session: graph = session.graph # Add hooks for the cell constructed above. with tf.variable_scope(component.name, reuse=True): runtime_support.add_hooks(component, cell) # Get the hook containing the wire-format proto. cell_wire_format = graph.get_tensor_by_name( '{}/EXPORT/CellSubgraphSpec:0'.format(component.name)) # Check that the hook matches the cell. tf.global_variables_initializer().run() self.assertEqual(cell_wire_format.eval(), cell.SerializeToString())
def __init__(self, master, component_spec, attr_defaults=None): """Initializes the ComponentBuilder from specifications. Args: master: dragnn.MasterBuilder object. component_spec: dragnn.ComponentSpec proto to be built. attr_defaults: Optional dict of component attribute defaults. If not provided or if empty, attributes are not extracted. """ self.master = master self.num_actions = component_spec.num_actions self.name = component_spec.name self.spec = component_spec self.moving_average = None # Determine if this component should apply self-normalization. self.eligible_for_self_norm = ( not self.master.hyperparams.self_norm_components_filter or self.name in self.master.hyperparams.self_norm_components_filter.split(',')) # Extract component attributes before make_network(), so the network unit # can access them. self._attrs = {} global_attr_defaults = { 'locally_normalize': False, 'output_as_probabilities': False } if attr_defaults: global_attr_defaults.update(attr_defaults) self._attrs = network_units.get_attrs_with_defaults( self.spec.component_builder.parameters, global_attr_defaults) do_local_norm = self._attrs['locally_normalize'] self._output_as_probabilities = self._attrs['output_as_probabilities'] with tf.variable_scope(self.name): self.training_beam_size = tf.constant(self.spec.training_beam_size, name='TrainingBeamSize') self.inference_beam_size = tf.constant( self.spec.inference_beam_size, name='InferenceBeamSize') self.locally_normalize = tf.constant(do_local_norm, name='LocallyNormalize') self._step = tf.get_variable('step', [], initializer=tf.zeros_initializer(), dtype=tf.int32) self._total = tf.get_variable('total', [], initializer=tf.zeros_initializer(), dtype=tf.int32) # Construct network variables. self.network = self.make_network(self.spec.network_unit) # Construct moving average. if self.master.hyperparams.use_moving_average: self.moving_average = tf.train.ExponentialMovingAverage( decay=self.master.hyperparams.average_weight, num_updates=self._step) self.avg_ops = [self.moving_average.apply(self.network.params)] # Used to export the cell; see add_cell_input() and add_cell_output(). self._cell_subgraph_spec = export_pb2.CellSubgraphSpec()
def testAddDerivedParamHooks(self): component = MockComponent() derived_name = 'derived' with self.test_session() as session: graph = session.graph # Add hooks. with tf.variable_scope(component.name, reuse=True): runtime_support.add_hooks(component, export_pb2.CellSubgraphSpec()) session.run(tf.global_variables_initializer()) # Get hooks for the derived vector. vector = graph.get_tensor_by_name('derived/vector:0') self.assertEqual(vector.shape, (3,)) # Get the hooks for the derived variable. matrix = graph.get_tensor_by_name( '{}/{}/matrix/blocked32:0'.format(component.name, derived_name)) self.assertAllEqual(tf.shape(matrix).eval(), [4, 128, 32]) # Check the bfloat16 version. It should have the same shape. bfloat16_matrix = graph.get_tensor_by_name( '{}/{}/matrix/blocked32/bfloat16:0'.format(component.name, derived_name)) self.assertAllEqual(tf.shape(bfloat16_matrix).eval(), [4, 128, 32])
def testAddFixedHooks(self): component = MockComponent() fixed0 = component.spec.fixed_feature.add() fixed1 = component.spec.fixed_feature.add() fixed0.embedding_dim = -1 fixed1.embedding_dim = 32 fixed0.vocabulary_size = 100 fixed1.vocabulary_size = 1000 fixed0_matrix_name = network_units.fixed_embeddings_name(0) fixed1_matrix_name = network_units.fixed_embeddings_name(1) with self.test_session() as session: graph = session.graph # Create fixed embedding matrices. Only channel 1 uses one. with tf.variable_scope(component.name): tf.get_variable( fixed1_matrix_name, shape=[1000 + 1, 32], dtype=tf.float32) # Add hooks. This should ignore channel 0 and add hooks for channel 1. with tf.variable_scope(component.name, reuse=True): runtime_support.add_hooks(component, export_pb2.CellSubgraphSpec()) # Check that no hooks were added for channel 0. with self.assertRaises(KeyError): graph.get_tensor_by_name( '{}/{}/trimmed:0'.format(component.name, fixed0_matrix_name)) # Get the hooks added for channel 1. trimmed = graph.get_tensor_by_name( '{}/{}/trimmed:0'.format(component.name, fixed1_matrix_name)) # Check dimensions of the hooks. tf.global_variables_initializer().run() self.assertAllEqual(tf.shape(trimmed).eval(), [1000, 32])
def testModelExportWithAveragesAndHooks(self): # Get the master spec and params for this graph. master_spec = self.LoadSpec('ud-hungarian.master-spec') params_path = os.path.join( test_flags.source_root(), 'dragnn/python/testdata' '/ud-hungarian.params') # Export the graph via SavedModel. (Here, we maintain a handle to the graph # for comparison, but that's usually not necessary.) Note that the export # path must not already exist. export_path = os.path.join(test_flags.temp_dir(), 'export2') dragnn_model_saver_lib.clean_output_paths(export_path) saver_graph = tf.Graph() shortened_to_original = dragnn_model_saver_lib.shorten_resource_paths( master_spec) dragnn_model_saver_lib.export_master_spec(master_spec, saver_graph) dragnn_model_saver_lib.export_to_graph(master_spec, params_path, export_path, saver_graph, export_moving_averages=True, build_runtime_graph=True) # Export the assets as well. dragnn_model_saver_lib.export_assets(master_spec, shortened_to_original, export_path) # Validate that the assets are all in the exported directory. path_set = self.ValidateAssetExistence(master_spec, export_path) # This master-spec has 4 unique assets. If there are more, we have not # uniquified the assets properly. self.assertEqual(len(path_set), 4) # Restore the graph from the checkpoint into a new Graph object. restored_graph = tf.Graph() restoration_config = tf.ConfigProto(log_device_placement=False, intra_op_parallelism_threads=10, inter_op_parallelism_threads=10) with tf.Session(graph=restored_graph, config=restoration_config) as sess: tf.saved_model.loader.load(sess, [tf.saved_model.tag_constants.SERVING], export_path) averaged_hook_name, non_averaged_hook_name, cell_subgraph_hook_name = ( self.GetHookNodeNames(master_spec)) # Check that an averaged runtime hook node exists. restored_graph.get_operation_by_name(averaged_hook_name) # Check that the non-averaged version does not exist. with self.assertRaises(KeyError): restored_graph.get_operation_by_name(non_averaged_hook_name) # Load the cell subgraph. cell_subgraph_bytes = restored_graph.get_tensor_by_name( cell_subgraph_hook_name + ':0') cell_subgraph_bytes = cell_subgraph_bytes.eval( feed_dict={'annotation/ComputeSession/InputBatch:0': []}) cell_subgraph_spec = export_pb2.CellSubgraphSpec() cell_subgraph_spec.ParseFromString(cell_subgraph_bytes) tf.logging.info('cell_subgraph_spec = %s', cell_subgraph_spec) # Sanity check inputs. for cell_input in cell_subgraph_spec.input: self.assertGreater(len(cell_input.name), 0) self.assertGreater(len(cell_input.tensor), 0) self.assertNotEqual( cell_input.type, export_pb2.CellSubgraphSpec.Input.TYPE_UNKNOWN) restored_graph.get_tensor_by_name( cell_input.tensor) # shouldn't raise # Sanity check outputs. for cell_output in cell_subgraph_spec.output: self.assertGreater(len(cell_output.name), 0) self.assertGreater(len(cell_output.tensor), 0) restored_graph.get_tensor_by_name( cell_output.tensor) # shouldn't raise # GetHookNames() finds a component with a fixed feature, so at least the # first feature ID should exist. self.assertTrue( any(cell_input.name == 'fixed_channel_0_index_0_ids' for cell_input in cell_subgraph_spec.input)) # Most dynamic components produce a logits layer. self.assertTrue( any(cell_output.name == 'logits' for cell_output in cell_subgraph_spec.output))
def testAddLinkedHooks(self): component = MockComponent() link0 = component.spec.linked_feature.add() link1 = component.spec.linked_feature.add() link0.embedding_dim = -1 # direct link link1.embedding_dim = 32 # transformed link link0_matrix_name = network_units.linked_embeddings_name(0) link1_matrix_name = network_units.linked_embeddings_name(1) with self.test_session() as session: graph = session.graph # Create linked embedding matrices. Only channel 1 uses one. with tf.variable_scope(component.name): tf.get_variable(link1_matrix_name, shape=[64 + 1, 32], dtype=tf.float32) # Add hooks. This should ignore channel 0 and add hooks for channel 1. with tf.variable_scope(component.name, reuse=True): runtime_support.add_hooks(component, export_pb2.CellSubgraphSpec()) # Check that no hooks were added for channel 0. with self.assertRaises(KeyError): graph.get_tensor_by_name( '{}/{}/weights:0'.format(component.name, link0_matrix_name)) with self.assertRaises(KeyError): graph.get_tensor_by_name('{}/{}/weights/transposed:0'.format( component.name, link0_matrix_name)) with self.assertRaises(KeyError): graph.get_tensor_by_name('{}/{}/weights/transposed/shape:0'.format( component.name, link0_matrix_name)) with self.assertRaises(KeyError): graph.get_tensor_by_name('{}/{}/weights/transposed/blocked32:0'.format( component.name, link0_matrix_name)) with self.assertRaises(KeyError): graph.get_tensor_by_name('{}/{}/weights/transposed/blocked48:0'.format( component.name, link0_matrix_name)) with self.assertRaises(KeyError): graph.get_tensor_by_name( '{}/{}/out_of_bounds:0'.format(component.name, link0_matrix_name)) # Get the hooks added for channel 1. weights = graph.get_tensor_by_name( '{}/{}/weights:0'.format(component.name, link1_matrix_name)) transposed = graph.get_tensor_by_name('{}/{}/weights/transposed:0'.format( component.name, link1_matrix_name)) transposed_shape = graph.get_tensor_by_name( '{}/{}/weights/transposed/shape:0'.format(component.name, link1_matrix_name)) transposed32 = graph.get_tensor_by_name( '{}/{}/weights/transposed/blocked32:0'.format(component.name, link1_matrix_name)) transposed48 = graph.get_tensor_by_name( '{}/{}/weights/transposed/blocked48:0'.format(component.name, link1_matrix_name)) out_of_bounds = graph.get_tensor_by_name( '{}/{}/out_of_bounds:0'.format(component.name, link1_matrix_name)) # Check dimensions of the hooks. tf.global_variables_initializer().run() self.assertAllEqual(tf.shape(weights).eval(), [64, 32]) self.assertAllEqual(tf.shape(transposed).eval(), [32, 64]) self.assertAllEqual(transposed_shape.eval(), [32, 64]) self.assertAllEqual(tf.shape(transposed32).eval(), [2, 32, 32]) self.assertAllEqual(tf.shape(transposed48).eval(), [2, 32, 48]) self.assertAllEqual(tf.shape(out_of_bounds).eval(), [1, 32])
def testAddParamsHooks(self): component = MockComponent() rank2_name = 'rank2' rank3_name = 'rank3' with self.test_session() as session: graph = session.graph # Add hooks. This should add hooks for all rank-2 params. with tf.variable_scope(component.name, reuse=True): runtime_support.add_hooks(component, export_pb2.CellSubgraphSpec()) # Check that no hooks were added for the rank-3 params. with self.assertRaises(KeyError): graph.get_tensor_by_name( '{}/{}/matrix:0'.format(component.name, rank3_name)) with self.assertRaises(KeyError): graph.get_tensor_by_name( '{}/{}/transposed:0'.format(component.name, rank3_name)) with self.assertRaises(KeyError): graph.get_tensor_by_name( '{}/{}/matrix/blocked32:0'.format(component.name, rank3_name)) with self.assertRaises(KeyError): graph.get_tensor_by_name( '{}/{}/matrix/blocked48:0'.format(component.name, rank3_name)) with self.assertRaises(KeyError): graph.get_tensor_by_name( '{}/{}/transposed/blocked32:0'.format(component.name, rank3_name)) with self.assertRaises(KeyError): graph.get_tensor_by_name( '{}/{}/transposed/blocked48:0'.format(component.name, rank3_name)) with self.assertRaises(KeyError): graph.get_tensor_by_name( '{}/{}/matrix/shape:0'.format(component.name, rank3_name)) with self.assertRaises(KeyError): graph.get_tensor_by_name( '{}/{}/transposed/shape:0'.format(component.name, rank3_name)) # Get the hooks added for each variable. matrix = graph.get_tensor_by_name( '{}/{}/matrix:0'.format(component.name, rank2_name)) transposed = graph.get_tensor_by_name( '{}/{}/transposed:0'.format(component.name, rank2_name)) matrix32 = graph.get_tensor_by_name( '{}/{}/matrix/blocked32:0'.format(component.name, rank2_name)) matrix48 = graph.get_tensor_by_name( '{}/{}/matrix/blocked48:0'.format(component.name, rank2_name)) transposed32 = graph.get_tensor_by_name( '{}/{}/transposed/blocked32:0'.format(component.name, rank2_name)) transposed48 = graph.get_tensor_by_name( '{}/{}/transposed/blocked48:0'.format(component.name, rank2_name)) matrix_shape = graph.get_tensor_by_name( '{}/{}/matrix/shape:0'.format(component.name, rank2_name)) transposed_shape = graph.get_tensor_by_name( '{}/{}/transposed/shape:0'.format(component.name, rank2_name)) # Check dimensions of the hooks. tf.global_variables_initializer().run() self.assertAllEqual(tf.shape(matrix).eval(), [64, 127]) self.assertAllEqual(tf.shape(transposed).eval(), [127, 64]) self.assertAllEqual(matrix_shape.eval(), [64, 127]) self.assertAllEqual(transposed_shape.eval(), [127, 64]) self.assertAllEqual(tf.shape(matrix32).eval(), [4, 64, 32]) self.assertAllEqual(tf.shape(matrix48).eval(), [3, 64, 48]) self.assertAllEqual(tf.shape(transposed32).eval(), [2, 127, 32]) self.assertAllEqual(tf.shape(transposed48).eval(), [2, 127, 48])