def create(self,
fixed_embeddings,
linked_embeddings,
context_tensor_arrays,
attention_tensor,
during_training,
stride=None):
"""Forwards the lengths and scores."""
check.NotNone(stride, 'MstSolverNetwork requires stride')
lengths = network_units.lookup_named_tensor('lengths',
linked_embeddings)
lengths_b = tf.to_int32(tf.squeeze(lengths.tensor, [1]))
scores = network_units.lookup_named_tensor('scores', linked_embeddings)
scores_bnxn = scores.tensor
max_length = tf.shape(scores_bnxn)[1]
scores_bxnxn = tf.reshape(scores_bnxn,
[stride, max_length, max_length])
_, argmax_sources_bxn = mst_ops.maximum_spanning_tree(
forest=self._attrs['forest'],
num_nodes=lengths_b,
scores=scores_bxnxn)
argmax_sources_bn = tf.reshape(argmax_sources_bxn, [-1])
arcs_bnxn = tf.one_hot(argmax_sources_bn, max_length, dtype=tf.float32)
return [lengths_b, scores_bxnxn, scores_bnxn, arcs_bnxn]
def create(self,
fixed_embeddings,
linked_embeddings,
context_tensor_arrays,
attention_tensor,
during_training,
stride=None):
"""Requires |stride|; otherwise see base class."""
del context_tensor_arrays, attention_tensor
if stride is None:
raise RuntimeError(
"PairwiseBilinearLabelNetwork needs 'stride' and must "
"be called in a bulk component.")
sources = network_units.lookup_named_tensor('sources',
linked_embeddings)
sources_tensor = tf.reshape(sources.tensor,
[stride, -1, self._source_dim])
targets = network_units.lookup_named_tensor('targets',
linked_embeddings)
targets_tensor = tf.reshape(targets.tensor,
[stride, -1, self._target_dim])
# Dimensions: source_dim x num_labels x target_dim
bilinear_params = self._component.get_variable('bilinear')
# Ensures that num_steps is the same for both inputs
num_steps = tf.shape(sources_tensor)[1]
with tf.control_dependencies([
tf.assert_equal(num_steps,
tf.shape(targets_tensor)[1],
name='num_steps_mismatch')
]):
# Dimensions:
# (batch_size*num_steps x source_dim) *
# (source_dim x num_labels*target_dim)
# = (batch_size*num_steps x num_labels*target_dim)
lin = tf.matmul(
tf.reshape(sources_tensor, [-1, self._source_dim]),
tf.reshape(bilinear_params, [self._source_dim, -1]))
# (batch_size x num_steps*num_labels x target_dim) *
# (batch_size x num_steps x target_dim)^T
# = (batch_size x num_steps*num_labels x num_steps)
bilin = tf.matmul(tf.reshape(
lin, [-1, num_steps * self._num_labels, self._target_dim]),
targets_tensor,
transpose_b=True)
# (batch_size x num_steps*num_labels x num_steps) ->
# (batch_size x num_steps x num_steps*num_labels)
scores = tf.transpose(bilin, [0, 2, 1])
return [
tf.reshape(scores, [-1, num_steps * self._num_labels],
name='reshape_activations')
]
def create(self,
fixed_embeddings,
linked_embeddings,
context_tensor_arrays,
attention_tensor,
during_training,
stride=None):
"""Requires |stride|; otherwise see base class."""
check.NotNone(
stride,
'BiaffineDigraphNetwork requires "stride" and must be called '
'in the bulk feature extractor component.')
# TODO(googleuser): Add dropout during training.
del during_training
# Retrieve (possibly averaged) weights.
weights_arc = self._component.get_variable('weights_arc')
weights_source = self._component.get_variable('weights_source')
root = self._component.get_variable('root')
# Extract the source and target token activations. Use |stride| to collapse
# batch and beam into a single dimension.
sources = network_units.lookup_named_tensor('sources',
linked_embeddings)
targets = network_units.lookup_named_tensor('targets',
linked_embeddings)
source_tokens_bxnxs = tf.reshape(sources.tensor,
[stride, -1, self._source_dim])
target_tokens_bxnxt = tf.reshape(targets.tensor,
[stride, -1, self._target_dim])
num_tokens = tf.shape(source_tokens_bxnxs)[1]
# Compute the arc, source, and root potentials.
arcs_bxnxn = digraph_ops.ArcPotentialsFromTokens(
source_tokens_bxnxs, target_tokens_bxnxt, weights_arc)
sources_bxnxn = digraph_ops.ArcSourcePotentialsFromTokens(
source_tokens_bxnxs, weights_source)
roots_bxn = digraph_ops.RootPotentialsFromTokens(
root, target_tokens_bxnxt, weights_arc, weights_source)
# Combine them into a single matrix with the roots on the diagonal.
adjacency_bxnxn = digraph_ops.CombineArcAndRootPotentials(
arcs_bxnxn + sources_bxnxn, roots_bxn)
# The adjacency matrix currently has sources on rows and targets on columns,
# but we want targets on rows so that maximizing within a row corresponds to
# selecting sources for a given target.
adjacency_bxnxn = tf.matrix_transpose(adjacency_bxnxn)
return [tf.reshape(adjacency_bxnxn, [-1, num_tokens])]
def create(self,
fixed_embeddings,
linked_embeddings,
context_tensor_arrays,
attention_tensor,
during_training,
stride=None):
"""Requires |stride|; otherwise see base class."""
check.NotNone(stride,
'BiaffineDigraphNetwork requires "stride" and must be called '
'in the bulk feature extractor component.')
# TODO(googleuser): Add dropout during training.
del during_training
# Retrieve (possibly averaged) weights.
weights_arc = self._component.get_variable('weights_arc')
weights_source = self._component.get_variable('weights_source')
root = self._component.get_variable('root')
# Extract the source and target token activations. Use |stride| to collapse
# batch and beam into a single dimension.
sources = network_units.lookup_named_tensor('sources', linked_embeddings)
targets = network_units.lookup_named_tensor('targets', linked_embeddings)
source_tokens_bxnxs = tf.reshape(sources.tensor,
[stride, -1, self._source_dim])
target_tokens_bxnxt = tf.reshape(targets.tensor,
[stride, -1, self._target_dim])
num_tokens = tf.shape(source_tokens_bxnxs)[1]
# Compute the arc, source, and root potentials.
arcs_bxnxn = digraph_ops.ArcPotentialsFromTokens(
source_tokens_bxnxs, target_tokens_bxnxt, weights_arc)
sources_bxnxn = digraph_ops.ArcSourcePotentialsFromTokens(
source_tokens_bxnxs, weights_source)
roots_bxn = digraph_ops.RootPotentialsFromTokens(
root, target_tokens_bxnxt, weights_arc, weights_source)
# Combine them into a single matrix with the roots on the diagonal.
adjacency_bxnxn = digraph_ops.CombineArcAndRootPotentials(
arcs_bxnxn + sources_bxnxn, roots_bxn)
# The adjacency matrix currently has sources on rows and targets on columns,
# but we want targets on rows so that maximizing within a row corresponds to
# selecting sources for a given target.
adjacency_bxnxn = tf.matrix_transpose(adjacency_bxnxn)
return [tf.reshape(adjacency_bxnxn, [-1, num_tokens])]
def create(self,
fixed_embeddings,
linked_embeddings,
context_tensor_arrays,
attention_tensor,
during_training,
stride=None):
"""Requires |stride|; otherwise see base class."""
del context_tensor_arrays, attention_tensor
if stride is None:
raise RuntimeError("PairwiseBilinearLabelNetwork needs 'stride' and must "
"be called in a bulk component.")
sources = network_units.lookup_named_tensor('sources', linked_embeddings)
sources_tensor = tf.reshape(sources.tensor, [stride, -1, self._source_dim])
targets = network_units.lookup_named_tensor('targets', linked_embeddings)
targets_tensor = tf.reshape(targets.tensor, [stride, -1, self._target_dim])
# Dimensions: source_dim x num_labels x target_dim
bilinear_params = self._component.get_variable('bilinear')
# Ensures that num_steps is the same for both inputs
num_steps = tf.shape(sources_tensor)[1]
with tf.control_dependencies([tf.assert_equal(num_steps,
tf.shape(targets_tensor)[1],
name='num_steps_mismatch')]):
# Dimensions:
# (batch_size*num_steps x source_dim) *
# (source_dim x num_labels*target_dim)
# = (batch_size*num_steps x num_labels*target_dim)
lin = tf.matmul(tf.reshape(sources_tensor, [-1, self._source_dim]),
tf.reshape(bilinear_params, [self._source_dim, -1]))
# (batch_size x num_steps*num_labels x target_dim) *
# (batch_size x num_steps x target_dim)^T
# = (batch_size x num_steps*num_labels x num_steps)
bilin = tf.matmul(
tf.reshape(lin, [-1, num_steps*self._num_labels, self._target_dim]),
targets_tensor, transpose_b=True)
# (batch_size x num_steps*num_labels x num_steps) ->
# (batch_size x num_steps x num_steps*num_labels)
scores = tf.transpose(bilin, [0, 2, 1])
return [tf.reshape(scores, [-1, num_steps*self._num_labels],
name='reshape_activations')]
def create(self,
fixed_embeddings,
linked_embeddings,
context_tensor_arrays,
attention_tensor,
during_training,
stride=None):
"""Requires |stride|; otherwise see base class."""
check.NotNone(
stride, 'BulkBiLSTMNetwork requires "stride" and must be called '
'in the bulk feature extractor component.')
# Flatten the lengths into a vector.
lengths = dragnn.lookup_named_tensor('lengths', linked_embeddings)
lengths_s = tf.squeeze(lengths.tensor, [1])
# Collect all other inputs into a batched tensor.
linked_embeddings = [
named_tensor for named_tensor in linked_embeddings
if named_tensor.name != 'lengths'
]
inputs_sxnxd = dragnn.get_input_tensor_with_stride(
fixed_embeddings, linked_embeddings, stride)
# Since get_input_tensor_with_stride() concatenates the input embeddings, it
# obscures the static activation dimension, which the RNN library requires.
# Restore it using set_shape(). Note that set_shape() merges into the known
# shape, so only specify the activation dimension.
inputs_sxnxd.set_shape(
[tf.Dimension(None),
tf.Dimension(None), self._input_dim])
initial_states_forward, initial_states_backward = (
self._create_initial_states(stride))
if during_training:
cells_forward = self._train_cells_forward
cells_backward = self._train_cells_backward
else:
cells_forward = self._inference_cells_forward
cells_backward = self._inference_cells_backward
def _bilstm_closure(scope):
"""Applies the bi-LSTM to the current inputs."""
outputs_sxnxd, _, _ = tf.contrib.rnn.stack_bidirectional_dynamic_rnn(
cells_forward,
cells_backward,
inputs_sxnxd,
initial_states_fw=initial_states_forward,
initial_states_bw=initial_states_backward,
sequence_length=lengths_s,
parallel_iterations=self._attrs['parallel_iterations'],
scope=scope)
return outputs_sxnxd
# Layer outputs are not batched; flatten out the batch dimension.
outputs_sxnxd = self._apply_with_captured_variables(_bilstm_closure)
outputs_snxd = tf.reshape(outputs_sxnxd, [-1, self._output_dim])
return self._append_base_layers([outputs_snxd])
def create(self,
fixed_embeddings,
linked_embeddings,
context_tensor_arrays,
attention_tensor,
during_training,
stride=None):
"""Requires |stride|; otherwise see base class."""
check.NotNone(
stride,
'BiaffineLabelNetwork requires "stride" and must be called '
'in the bulk feature extractor component.')
# TODO(googleuser): Add dropout during training.
del during_training
# Retrieve (possibly averaged) weights.
weights_pair = self._component.get_variable('weights_pair')
weights_source = self._component.get_variable('weights_source')
weights_target = self._component.get_variable('weights_target')
biases = self._component.get_variable('biases')
# Extract and shape the source and target token activations. Use |stride|
# to collapse batch and beam into a single dimension.
sources = network_units.lookup_named_tensor('sources',
linked_embeddings)
targets = network_units.lookup_named_tensor('targets',
linked_embeddings)
sources_bxnxs = tf.reshape(sources.tensor,
[stride, -1, self._source_dim])
targets_bxnxt = tf.reshape(targets.tensor,
[stride, -1, self._target_dim])
# Compute the pair, source, and target potentials.
pairs_bxnxl = digraph_ops.LabelPotentialsFromTokenPairs(
sources_bxnxs, targets_bxnxt, weights_pair)
sources_bxnxl = digraph_ops.LabelPotentialsFromTokens(
sources_bxnxs, weights_source)
targets_bxnxl = digraph_ops.LabelPotentialsFromTokens(
targets_bxnxt, weights_target)
# Combine them with the biases.
labels_bxnxl = pairs_bxnxl + sources_bxnxl + targets_bxnxl + biases
# Flatten out the batch dimension.
return [tf.reshape(labels_bxnxl, [-1, self._num_labels])]
def create(self,
fixed_embeddings,
linked_embeddings,
context_tensor_arrays,
attention_tensor,
during_training,
stride=None):
"""Requires |stride|; otherwise see base class."""
check.NotNone(stride,
'BulkBiLSTMNetwork requires "stride" and must be called '
'in the bulk feature extractor component.')
# Flatten the lengths into a vector.
lengths = dragnn.lookup_named_tensor('lengths', linked_embeddings)
lengths_s = tf.squeeze(lengths.tensor, [1])
# Collect all other inputs into a batched tensor.
linked_embeddings = [
named_tensor for named_tensor in linked_embeddings
if named_tensor.name != 'lengths'
]
inputs_sxnxd = dragnn.get_input_tensor_with_stride(
fixed_embeddings, linked_embeddings, stride)
# Since get_input_tensor_with_stride() concatenates the input embeddings, it
# obscures the static activation dimension, which the RNN library requires.
# Restore it using set_shape(). Note that set_shape() merges into the known
# shape, so only specify the activation dimension.
inputs_sxnxd.set_shape(
[tf.Dimension(None), tf.Dimension(None), self._input_dim])
initial_states_forward, initial_states_backward = (
self._create_initial_states(stride))
if during_training:
cells_forward = self._train_cells_forward
cells_backward = self._train_cells_backward
else:
cells_forward = self._inference_cells_forward
cells_backward = self._inference_cells_backward
def _bilstm_closure(scope):
"""Applies the bi-LSTM to the current inputs."""
outputs_sxnxd, _, _ = tf.contrib.rnn.stack_bidirectional_dynamic_rnn(
cells_forward,
cells_backward,
inputs_sxnxd,
initial_states_fw=initial_states_forward,
initial_states_bw=initial_states_backward,
sequence_length=lengths_s,
parallel_iterations=self._attrs['parallel_iterations'],
scope=scope)
return outputs_sxnxd
# Layer outputs are not batched; flatten out the batch dimension.
outputs_sxnxd = self._apply_with_captured_variables(_bilstm_closure)
outputs_snxd = tf.reshape(outputs_sxnxd, [-1, self._output_dim])
return self._append_base_layers([outputs_snxd])
def create(self,
fixed_embeddings,
linked_embeddings,
context_tensor_arrays,
attention_tensor,
during_training,
stride=None):
"""Requires |stride|; otherwise see base class."""
check.NotNone(stride,
'BiaffineLabelNetwork requires "stride" and must be called '
'in the bulk feature extractor component.')
# TODO(googleuser): Add dropout during training.
del during_training
# Retrieve (possibly averaged) weights.
weights_pair = self._component.get_variable('weights_pair')
weights_source = self._component.get_variable('weights_source')
weights_target = self._component.get_variable('weights_target')
biases = self._component.get_variable('biases')
# Extract and shape the source and target token activations. Use |stride|
# to collapse batch and beam into a single dimension.
sources = network_units.lookup_named_tensor('sources', linked_embeddings)
targets = network_units.lookup_named_tensor('targets', linked_embeddings)
sources_bxnxs = tf.reshape(sources.tensor, [stride, -1, self._source_dim])
targets_bxnxt = tf.reshape(targets.tensor, [stride, -1, self._target_dim])
# Compute the pair, source, and target potentials.
pairs_bxnxl = digraph_ops.LabelPotentialsFromTokenPairs(sources_bxnxs,
targets_bxnxt,
weights_pair)
sources_bxnxl = digraph_ops.LabelPotentialsFromTokens(sources_bxnxs,
weights_source)
targets_bxnxl = digraph_ops.LabelPotentialsFromTokens(targets_bxnxt,
weights_target)
# Combine them with the biases.
labels_bxnxl = pairs_bxnxl + sources_bxnxl + targets_bxnxl + biases
# Flatten out the batch dimension.
return [tf.reshape(labels_bxnxl, [-1, self._num_labels])]
def create(self,
fixed_embeddings,
linked_embeddings,
context_tensor_arrays,
attention_tensor,
during_training,
stride=None):
"""Forwards the lengths and scores."""
check.NotNone(stride, 'MstSolverNetwork requires stride')
lengths = network_units.lookup_named_tensor('lengths', linked_embeddings)
lengths_b = tf.to_int32(tf.squeeze(lengths.tensor, [1]))
scores = network_units.lookup_named_tensor('scores', linked_embeddings)
scores_bnxn = scores.tensor
max_length = tf.shape(scores_bnxn)[1]
scores_bxnxn = tf.reshape(scores_bnxn, [stride, max_length, max_length])
_, argmax_sources_bxn = mst_ops.maximum_spanning_tree(
forest=self._attrs['forest'], num_nodes=lengths_b, scores=scores_bxnxn)
argmax_sources_bn = tf.reshape(argmax_sources_bxn, [-1])
arcs_bnxn = tf.one_hot(argmax_sources_bn, max_length, dtype=tf.float32)
return [lengths_b, scores_bxnxn, scores_bnxn, arcs_bnxn]
def create(self,
fixed_embeddings,
linked_embeddings,
context_tensor_arrays,
attention_tensor,
during_training,
stride=None):
"""Requires |stride|; otherwise see base class."""
del context_tensor_arrays, attention_tensor
if stride is None:
raise RuntimeError("TransformerEncoderNetwork needs 'stride' and must be "
"called in the bulk feature extractor component.")
lengths = network_units.lookup_named_tensor('lengths', linked_embeddings)
lengths_s = tf.to_int32(tf.squeeze(lengths.tensor, [1]))
num_steps = tf.reduce_max(lengths_s)
in_tensor = network_units.lookup_named_tensor('features', linked_embeddings)
input_tensor = tf.reshape(in_tensor.tensor, [stride, num_steps, -1])
if self._timing_signal:
input_tensor = add_timing_signal_1d(input_tensor)
# Adds a dimension for conv2d
input_tensor = tf.expand_dims(input_tensor, 1)
# For masking padding in attention
mask = compute_padding_mask(lengths_s)
conv = tf.nn.conv2d(input_tensor,
self._component.get_variable('init_proj'),
[1, 1, 1, 1], padding='SAME')
conv = tf.nn.bias_add(conv, self._component.get_variable('init_bias'))
for i in range(self._num_layers):
with tf.variable_scope('transform_%d' % i, reuse=True):
attn_weights = self._component.get_variable('attn_weights')
attn_combined = tf.nn.conv2d(conv,
attn_weights,
[1, 1, 1, 1],
padding='SAME')
attn_combined = tf.squeeze(attn_combined, 1)
# Splits combined projection into queries, keys, and values
queries, keys, values = tf.split(attn_combined,
[self._combined_filters]*3,
axis=2)
# Splits each of queries, keys, values into attention heads
queries = split_heads(queries, self._num_heads)
keys = split_heads(keys, self._num_heads)
values = split_heads(values, self._num_heads)
if self._scale_attn:
queries *= self._filter_size**-0.5
# Performs dot product attention and concatenates the resulting heads
attended = dot_product_attention(queries, keys, values,
self._attention_dropout, mask)
attended = combine_heads(attended)
# Projects combined heads
attended = tf.expand_dims(attended, 1)
proj = tf.nn.conv2d(attended,
self._component.get_variable('proj_weights'),
[1, 1, 1, 1],
padding='SAME')
# Residual connection between input and attended input
attn_layer_norm_params = None
if self._layer_norm_res:
attn_layer_norm_params = self._layer_norms['attn_layer_norm_%d' % i]
proj_res = residual(conv, proj, self._residual_dropout,
attn_layer_norm_params)
# Feed forward
with tf.variable_scope('mlp'):
ff = mlp(self._component, proj_res, self._mlp_dropout,
self._mlp_depth)
# Residual connection between attended input and feed forward layers
ff_layer_norm_params = None
if self._layer_norm_res:
ff_layer_norm_params = self._layer_norms['ff_layer_norm_%d' % i]
conv = residual(proj_res, ff, self._residual_dropout,
ff_layer_norm_params)
return [tf.reshape(conv, [-1, self._combined_filters],
name='reshape_activations')]
def create(self,
fixed_embeddings,
linked_embeddings,
context_tensor_arrays,
attention_tensor,
during_training,
stride=None):
"""Requires |stride|; otherwise see base class."""
del context_tensor_arrays, attention_tensor
if stride is None:
raise RuntimeError("TransformerEncoderNetwork needs 'stride' and must be "
"called in the bulk feature extractor component.")
lengths = network_units.lookup_named_tensor('lengths', linked_embeddings)
lengths_s = tf.to_int32(tf.squeeze(lengths.tensor, [1]))
num_steps = tf.reduce_max(lengths_s)
in_tensor = network_units.lookup_named_tensor('features', linked_embeddings)
input_tensor = tf.reshape(in_tensor.tensor, [stride, num_steps, -1])
if self._timing_signal:
input_tensor = add_timing_signal_1d(input_tensor)
# Adds a dimension for conv2d
input_tensor = tf.expand_dims(input_tensor, 1)
# For masking padding in attention
mask = compute_padding_mask(lengths_s)
conv = tf.nn.conv2d(input_tensor,
self._component.get_variable('init_proj'),
[1, 1, 1, 1], padding='SAME')
conv = tf.nn.bias_add(conv, self._component.get_variable('init_bias'))
for i in range(self._num_layers):
with tf.variable_scope('transform_%d' % i, reuse=True):
attn_weights = self._component.get_variable('attn_weights')
attn_combined = tf.nn.conv2d(conv,
attn_weights,
[1, 1, 1, 1],
padding='SAME')
attn_combined = tf.squeeze(attn_combined, 1)
# Splits combined projection into queries, keys, and values
queries, keys, values = tf.split(attn_combined,
[self._combined_filters]*3,
axis=2)
# Splits each of queries, keys, values into attention heads
queries = split_heads(queries, self._num_heads)
keys = split_heads(keys, self._num_heads)
values = split_heads(values, self._num_heads)
if self._scale_attn:
queries *= self._filter_size**-0.5
# Performs dot product attention and concatenates the resulting heads
attended = dot_product_attention(queries, keys, values,
self._attention_dropout, mask)
attended = combine_heads(attended)
# Projects combined heads
attended = tf.expand_dims(attended, 1)
proj = tf.nn.conv2d(attended,
self._component.get_variable('proj_weights'),
[1, 1, 1, 1],
padding='SAME')
# Residual connection between input and attended input
attn_layer_norm_params = None
if self._layer_norm_res:
attn_layer_norm_params = self._layer_norms['attn_layer_norm_%d' % i]
proj_res = residual(conv, proj, self._residual_dropout,
attn_layer_norm_params)
# Feed forward
with tf.variable_scope('mlp'):
ff = mlp(self._component, proj_res, self._mlp_dropout,
self._mlp_depth)
# Residual connection between attended input and feed forward layers
ff_layer_norm_params = None
if self._layer_norm_res:
ff_layer_norm_params = self._layer_norms['ff_layer_norm_%d' % i]
conv = residual(proj_res, ff, self._residual_dropout,
ff_layer_norm_params)
return [tf.reshape(conv, [-1, self._combined_filters],
name='reshape_activations')]