def testConstantPadding(self):
with tf.Graph().as_default(), self.test_session():
with tf.variable_scope('test_scope'):
network = network_units.GatherNetwork(self._component)
# Construct a batch of two items with 3 and 2 steps, respectively.
indices = tf.constant([[1], [2], [0], # item 1
[-1], [0], [-1]], # item 2
dtype=tf.int64)
features = tf.constant([[1.0, 1.5], [2.0, 2.5], [3.0, 3.5], # item 1
[4.0, 4.5], [5.0, 5.5], [6.0, 6.5]], # item 2
dtype=tf.float32)
fixed_embeddings = []
linked_embeddings = [
network_units.NamedTensor(indices, 'indices', 1),
network_units.NamedTensor(features, 'features', 2)
]
with tf.variable_scope('test_scope', reuse=True):
outputs = network.create(fixed_embeddings, linked_embeddings, None,
None, True, 2)
gathered = outputs[0]
# Zeros will be substituted for index -1.
self.assertAllEqual(gathered.eval(),
[[2.0, 2.5], # gathered from 1
[3.0, 3.5], # gathered from 2
[1.0, 1.5], # gathered from 0
[0.0, 0.0], # gathered from -1
[4.0, 4.5], # gathered from 0
[0.0, 0.0]]) # gathered from -1
def testCreate(self):
with self.test_session():
master = MockMaster()
component = MockComponent(master, master.spec.component[0])
component.network = mst_units.MstSolverNetwork(component)
stride = 1
lengths = tf.constant([[3]], dtype=tf.int64)
scores = tf.constant(
[[1.0, 0.5, 0.5], [2.0, 0.5, 0.5], [0.5, 3.0, 0.5]],
dtype=tf.float32) # pyformat: disable
linked_embeddings = [
network_units.NamedTensor(lengths, 'lengths'),
network_units.NamedTensor(scores, 'scores')
]
network_tensors = component.network.create([], linked_embeddings,
[], None, False, stride)
self.assertAllEqual(network_tensors[0].eval(), [3])
self.assertAllEqual(network_tensors[1].eval(),
[[[1.0, 0.5, 0.5], [2.0, 0.5, 0.5],
[0.5, 3.0, 0.5]]]) # pyformat: disable
self.assertAllEqual(network_tensors[2].eval(),
[[1.0, 0.5, 0.5], [2.0, 0.5, 0.5],
[0.5, 3.0, 0.5]]) # pyformat: disable
self.assertAllEqual(network_tensors[3].eval(),
[[1.0, 0.0, 0.0], [1.0, 0.0, 0.0],
[0.0, 1.0, 0.0]]) # pyformat: disable
def testCanCreate(self):
"""Tests that create() works on a good spec."""
with tf.Graph().as_default(), self.test_session():
master = MockMaster()
component = MockComponent(master, _make_biaffine_spec())
with tf.variable_scope(component.name, reuse=None):
component.network = biaffine_units.BiaffineDigraphNetwork(
component)
with tf.variable_scope(component.name, reuse=True):
sources = network_units.NamedTensor(
tf.zeros([_BATCH_SIZE * _NUM_TOKENS, _TOKEN_DIM]),
'sources')
targets = network_units.NamedTensor(
tf.zeros([_BATCH_SIZE * _NUM_TOKENS, _TOKEN_DIM]),
'targets')
# No assertions on the result, just don't crash.
component.network.create(fixed_embeddings=[],
linked_embeddings=[sources, targets],
context_tensor_arrays=None,
attention_tensor=None,
during_training=True,
stride=_BATCH_SIZE)
def testGetBulkPredictions(self):
with self.test_session():
master = MockMaster()
component = MockComponent(master, master.spec.component[0])
component.network = mst_units.MstSolverNetwork(component)
stride = 2
lengths = tf.constant([[2], [3]], dtype=tf.int64)
pad = -12345.6
scores = tf.constant(
[[1.0, 2.0, pad], [1.8, 2.0, pad], [pad, pad, pad],
[3.8, 4.0, 3.9], [3.9, 3.8, 4.0], [3.8, 0.9, 4.0]],
dtype=tf.float32) # pyformat: disable
linked_embeddings = [
network_units.NamedTensor(lengths, 'lengths'),
network_units.NamedTensor(scores, 'scores')
]
network_tensors = component.network.create([], linked_embeddings,
[], None, False, stride)
predictions = component.network.get_bulk_predictions(
stride, network_tensors)
self.assertAllEqual(
predictions.eval(),
[[0.0, 1.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 0.0],
[0.0, 1.0, 0.0], [0.0, 0.0, 1.0], [0.0, 0.0, 1.0]
]) # pyformat: disable
def testTrainablePadding(self):
self._component.spec.network_unit.parameters[
'trainable_padding'] = 'true'
with tf.Graph().as_default(), self.test_session():
with tf.variable_scope('test_scope'):
network = network_units.GatherNetwork(self._component)
# Construct a batch of two items with 3 and 2 steps, respectively.
indices = tf.constant(
[
[1],
[2],
[0], # item 1
[-1],
[0],
[-1]
], # item 2
dtype=tf.int64)
features = tf.constant(
[
[1.0, 1.5],
[2.0, 2.5],
[3.0, 3.5], # item 1
[4.0, 4.5],
[5.0, 5.5],
[6.0, 6.5]
], # item 2
dtype=tf.float32)
fixed_embeddings = []
linked_embeddings = [
network_units.NamedTensor(indices, 'indices', 1),
network_units.NamedTensor(features, 'features', 2)
]
with tf.variable_scope('test_scope', reuse=True):
outputs = network.create(fixed_embeddings, linked_embeddings,
None, None, True, 2)
gathered = outputs[0]
# Ensure that the padding variable is initialized.
tf.global_variables_initializer().run()
# Randomly-initialized padding will be substituted for index -1.
self.assertAllEqual(gathered[0].eval(),
[2.0, 2.5]) # gathered from 1
self.assertAllEqual(gathered[1].eval(),
[3.0, 3.5]) # gathered from 2
self.assertAllEqual(gathered[2].eval(),
[1.0, 1.5]) # gathered from 0
tf.logging.info('padding = %s',
gathered[3].eval()) # gathered from -1
self.assertAllEqual(gathered[4].eval(),
[4.0, 4.5]) # gathered from 0
tf.logging.info('padding = %s',
gathered[5].eval()) # gathered from -1
# Though random, the padding must identical.
self.assertAllEqual(gathered[3].eval(), gathered[5].eval())
def testComputeBulkLossCrf(self):
with self.test_session():
master = MockMaster()
component = MockComponent(master, master.spec.component[0])
component.spec.network_unit.parameters['loss'] = 'crf'
component.network = mst_units.MstSolverNetwork(component)
stride = 2
lengths = tf.constant([[2], [3]], dtype=tf.int64)
# These scores have 2.0 (in the log domain) on the gold arcs and 1.0
# elsewhere.
pad = -12345.6
one = math.log(1.0)
two = math.log(2.0)
scores = tf.constant(
[[one, two, pad], [one, two, pad], [pad, pad, pad],
[one, two, one], [one, one, two], [one, one, two]],
dtype=tf.float32) # pyformat: disable
gold = tf.constant([1, 1, -1, 1, 2, 2], tf.int32)
first_partition_function = (
2.0 * 2.0 + # 0 -> 1 (gold)
1.0 * 1.0) # 1 -> 0
first_loss = -math.log(2.0 * 2.0 / first_partition_function)
second_partition_function = (
2.0 * 2.0 * 2.0 + # 0 -> 1 -> 2 (gold)
1.0 * 1.0 * 1.0 + # 2 -> 1 -> 0
1.0 * 1.0 * 1.0 + # 0 -> 2 -> 1
2.0 * 1.0 * 1.0 + # 1 -> 2 -> 0
2.0 * 1.0 * 1.0 + # 1 -> 0 -> 2
2.0 * 1.0 * 1.0 + # 2 -> 0 -> 1
2.0 * 2.0 * 1.0 + # {0, 1} -> 2
2.0 * 1.0 * 1.0 + # {0, 2} -> 1
1.0 * 1.0 * 1.0) # {1, 2} -> 0
second_loss = -math.log(
2.0 * 2.0 * 2.0 / second_partition_function)
linked_embeddings = [
network_units.NamedTensor(lengths, 'lengths'),
network_units.NamedTensor(scores, 'scores')
]
network_tensors = component.network.create([], linked_embeddings,
[], None, False, stride)
cost, correct, total = component.network.compute_bulk_loss(
stride, network_tensors, gold)
self.assertAlmostEqual(cost.eval(), first_loss + second_loss)
self.assertEqual(correct.eval(), 2 + 3)
self.assertEqual(total.eval(), 2 + 3)
def fetch_linked_embedding(comp, network_states, feature_spec):
"""Looks up linked embeddings in other components.
Args:
comp: ComponentBuilder object with respect to which the feature is to be
fetched
network_states: dictionary of NetworkState objects
feature_spec: FeatureSpec proto for the linked feature to be looked up
Returns:
NamedTensor containing the linked feature tensor
Raises:
NotImplementedError: if a linked feature with source translator other than
'identity' is configured.
RuntimeError: if a recurrent linked feature is configured.
"""
if feature_spec.source_translator != 'identity':
raise NotImplementedError(feature_spec.source_translator)
if feature_spec.source_component == comp.name:
raise RuntimeError(
'Recurrent linked features are not supported in bulk extraction.')
tf.logging.info('[%s] Adding linked feature "%s"', comp.name,
feature_spec.name)
source = comp.master.lookup_component[feature_spec.source_component]
return network_units.NamedTensor(
network_states[source.name].activations[
feature_spec.source_layer].bulk_tensor,
feature_spec.name)
def fetch_fast_fixed_embeddings(comp, state):
"""Looks up fixed features with fast, non-differentiable, op.
Since BulkFixedEmbeddings is non-differentiable with respect to the
embeddings, the idea is to call this function only when the graph is
not being used for training.
Args:
comp: Component whose fixed features we wish to look up.
state: live MasterState object for the component.
Returns:
state handle: updated state handle to be used after this call
fixed_embeddings: list of NamedTensor objects
"""
_validate_embedded_fixed_features(comp)
num_channels = len(comp.spec.fixed_feature)
if not num_channels:
return state.handle, []
tf.logging.info('[%s] Adding %d fast fixed features', comp.name, num_channels)
state.handle, bulk_embeddings, _ = dragnn_ops.bulk_fixed_embeddings(
state.handle, [
comp.get_variable(network_units.fixed_embeddings_name(c))
for c in range(num_channels)
],
component=comp.name)
bulk_embeddings = network_units.NamedTensor(bulk_embeddings,
'bulk-%s-fixed-features' %
comp.name)
return state.handle, [bulk_embeddings]
def testComputeBulkLossM3n(self):
with self.test_session():
master = MockMaster()
component = MockComponent(master, master.spec.component[0])
component.spec.network_unit.parameters['loss'] = 'm3n'
component.network = mst_units.MstSolverNetwork(component)
stride = 2
lengths = tf.constant([[2], [3]], dtype=tf.int64)
# Note that these scores are large enough to overcome the +1 hamming loss
# terms in the M3N loss. Therefore, the score matrix determines the tree
# that is used to compute the M3N loss.
pad = -12345.6
scores = tf.constant(
[[0.5, 2.0, pad], [0.5, 2.0, pad], [pad, pad, pad],
[2.5, 4.0, 2.5], [2.5, 2.5, 4.0], [2.5, 2.5, 4.0]],
dtype=tf.float32) # pyformat: disable
# For the first tree, the gold and scores agree on one arc (that index 1
# is a root), and for the second tree, the gold and scores agree on none
# of the arcs. Therefore, we expect +1 and +3 for the first and second
# trees in the M3N loss.
gold = tf.constant([0, 1, -1, 0, 0, 1], tf.int32)
first_gold_score = 0.5 + 2.0
second_gold_score = 2.5 + 2.5 + 2.5
first_tree_correct = 1
second_tree_correct = 0
first_tree_loss = 2 * 2.0 + 2 - first_tree_correct - first_gold_score
second_tree_loss = 3 * 4.0 + 3 - second_tree_correct - second_gold_score
linked_embeddings = [
network_units.NamedTensor(lengths, 'lengths'),
network_units.NamedTensor(scores, 'scores')
]
network_tensors = component.network.create([], linked_embeddings,
[], None, False, stride)
cost, correct, total = component.network.compute_bulk_loss(
stride, network_tensors, gold)
self.assertEqual(cost.eval(), first_tree_loss + second_tree_loss)
self.assertEqual(correct.eval(),
first_tree_correct + second_tree_correct)
self.assertEqual(total.eval(), 2 + 3)
def fetch_dense_ragged_embeddings(comp, state):
"""Gets embeddings in RaggedTensor format."""
_validate_embedded_fixed_features(comp)
num_channels = len(comp.spec.fixed_feature)
if not num_channels:
return state.handle, []
tf.logging.info('[%s] Adding %d fast fixed features', comp.name, num_channels)
features = [
comp.get_variable(network_units.fixed_embeddings_name(c))
for c in range(num_channels)
]
state.handle, data, offsets = dragnn_ops.bulk_embed_dense_fixed_features(
state.handle, features, component=comp.name)
data = network_units.NamedTensor(data, 'dense-%s-data' % comp.name)
offsets = network_units.NamedTensor(offsets, 'dense-%s-offsets' % comp.name)
return state.handle, [data, offsets]
def fetch_fast_fixed_embeddings(comp,
state,
pad_to_batch=None,
pad_to_steps=None):
"""Looks up fixed features with fast, non-differentiable, op.
Since BulkFixedEmbeddings is non-differentiable with respect to the
embeddings, the idea is to call this function only when the graph is
not being used for training. If the function is being called with fixed step
and batch sizes, it will use the most efficient possible extractor.
Args:
comp: Component whose fixed features we wish to look up.
state: live MasterState object for the component.
pad_to_batch: Optional; the number of batch elements to pad to.
pad_to_steps: Optional; the number of steps to pad to.
Returns:
state handle: updated state handle to be used after this call
fixed_embeddings: list of NamedTensor objects
"""
_validate_embedded_fixed_features(comp)
num_channels = len(comp.spec.fixed_feature)
if not num_channels:
return state.handle, []
tf.logging.info('[%s] Adding %d fast fixed features', comp.name,
num_channels)
features = [
comp.get_variable(network_units.fixed_embeddings_name(c))
for c in range(num_channels)
]
if pad_to_batch is not None and pad_to_steps is not None:
# If we have fixed padding numbers, we can use 'bulk_embed_fixed_features',
# which is the fastest embedding extractor.
state.handle, bulk_embeddings, _ = dragnn_ops.bulk_embed_fixed_features(
state.handle,
features,
component=comp.name,
pad_to_batch=pad_to_batch,
pad_to_steps=pad_to_steps)
else:
state.handle, bulk_embeddings, _ = dragnn_ops.bulk_fixed_embeddings(
state.handle, features, component=comp.name)
bulk_embeddings = network_units.NamedTensor(
bulk_embeddings, 'bulk-%s-fixed-features' % comp.name)
return state.handle, [bulk_embeddings]
def fetch_differentiable_fixed_embeddings(comp, state, stride,
during_training):
"""Looks up fixed features with separate, differentiable, embedding lookup.
Args:
comp: Component whose fixed features we wish to look up.
state: live MasterState object for the component.
stride: Tensor containing current batch * beam size.
during_training: True if this is being called from a training code path.
This controls, e.g., the use of feature ID dropout.
Returns:
state handle: updated state handle to be used after this call
fixed_embeddings: list of NamedTensor objects
"""
_validate_embedded_fixed_features(comp)
num_channels = len(comp.spec.fixed_feature)
if not num_channels:
return state.handle, []
state.handle, indices, ids, weights, num_steps = (
dragnn_ops.bulk_fixed_features(state.handle,
component=comp.name,
num_channels=num_channels))
fixed_embeddings = []
for channel, feature_spec in enumerate(comp.spec.fixed_feature):
differentiable_or_constant = ('constant' if feature_spec.is_constant
else 'differentiable')
tf.logging.info('[%s] Adding %s fixed feature "%s"', comp.name,
differentiable_or_constant, feature_spec.name)
if during_training and feature_spec.dropout_id >= 0:
ids[channel], weights[
channel] = network_units.apply_feature_id_dropout(
ids[channel], weights[channel], feature_spec)
size = stride * num_steps * feature_spec.size
fixed_embedding = network_units.embedding_lookup(
comp.get_variable(network_units.fixed_embeddings_name(channel)),
indices[channel], ids[channel], weights[channel], size)
if feature_spec.is_constant:
fixed_embedding = tf.stop_gradient(fixed_embedding)
fixed_embeddings.append(
network_units.NamedTensor(fixed_embedding, feature_spec.name))
return state.handle, fixed_embeddings
def extract_fixed_feature_ids(comp, state, stride):
"""Extracts fixed feature IDs.
Args:
comp: Component whose fixed feature IDs we wish to extract.
state: Live MasterState object for the component.
stride: Tensor containing current batch * beam size.
Returns:
state handle: Updated state handle to be used after this call.
ids: List of [stride * num_steps, 1] feature IDs per channel. Missing IDs
(e.g., due to batch padding) are set to -1.
"""
num_channels = len(comp.spec.fixed_feature)
if not num_channels:
return state.handle, []
for feature_spec in comp.spec.fixed_feature:
check.Eq(feature_spec.size, 1, 'All features must have size=1')
check.Lt(feature_spec.embedding_dim, 0,
'All features must be non-embedded')
state.handle, indices, ids, _, num_steps = dragnn_ops.bulk_fixed_features(
state.handle, component=comp.name, num_channels=num_channels)
size = stride * num_steps
fixed_ids = []
for channel, feature_spec in enumerate(comp.spec.fixed_feature):
tf.logging.info('[%s] Adding fixed feature IDs "%s"', comp.name,
feature_spec.name)
# The +1 and -1 increments ensure that missing IDs default to -1.
#
# TODO(googleuser): This formula breaks if multiple IDs are extracted at some
# step. Try using tf.unique() to enforce the unique-IDS precondition.
sums = tf.unsorted_segment_sum(ids[channel] + 1, indices[channel],
size) - 1
sums = tf.expand_dims(sums, axis=1)
fixed_ids.append(
network_units.NamedTensor(sums, feature_spec.name, dim=1))
return state.handle, fixed_ids
def fixed_word_embeddings(self):
"""Helper for returning fixed embeddings, for 1 word feature."""
words_tensor = tf.constant([[1.0] * 32], dtype=tf.float32)
return [network_units.NamedTensor(words_tensor, 'words')]
