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