def mlp(feature, hparams, name="mlp"):
"""Multi layer perceptron with dropout and relu activation."""
with tf.variable_scope(name, "mlp", values=[feature]):
num_mlp_layers = hparams.num_mlp_layers
mlp_dim = hparams.mlp_dim
for _ in range(num_mlp_layers):
feature = common_layers.dense(feature, mlp_dim, activation=tf.nn.relu)
feature = tf.nn.dropout(feature, keep_prob=1.-hparams.dropout)
return feature
def body(self, features):
hp = self.hparams
# pylint: disable=eval-used
if hp.image_input_type == "image":
image_feat = vqa_layers.image_embedding(
features["inputs"],
model_fn=eval(hp.image_model_fn),
trainable=hp.train_resnet,
is_training=hp.mode == tf.estimator.ModeKeys.TRAIN)
else:
image_feat = features["inputs"]
image_feat = common_layers.flatten4d3d(image_feat)
image_feat = common_layers.dense(image_feat, hp.hidden_size)
utils.collect_named_outputs("norms", "image_feat_after_proj",
tf.norm(image_feat, axis=-1))
question = common_layers.flatten4d3d(features["question"])
utils.collect_named_outputs("norms", "question_embedding",
tf.norm(question, axis=-1))
(encoder_input, encoder_self_attention_bias,
encoder_decoder_attention_bias) = prepare_image_question_encoder(
image_feat, question, hp)
encoder_input = tf.nn.dropout(
encoder_input, keep_prob=1.-hp.layer_prepostprocess_dropout)
encoder_output, _ = recurrent_transformer_decoder(
encoder_input, None, encoder_self_attention_bias, None,
hp, name="encoder")
utils.collect_named_outputs(
"norms", "encoder_output", tf.norm(encoder_output, axis=-1))
# scale query by sqrt(hidden_size)
query = tf.get_variable("query", [hp.hidden_size]) * hp.hidden_size **0.5
query = tf.expand_dims(tf.expand_dims(query, axis=0), axis=0)
batch_size = common_layers.shape_list(encoder_input)[0]
query = tf.tile(query, [batch_size, 1, 1])
query = tf.nn.dropout(
query, keep_prob=1.-hp.layer_prepostprocess_dropout)
decoder_output, _ = recurrent_transformer_decoder(
query, encoder_output, None, encoder_decoder_attention_bias,
hp, name="decoder")
utils.collect_named_outputs("norms", "decoder_output",
tf.norm(decoder_output, axis=-1))
norm_tensors = utils.convert_collection_to_dict("norms")
vqa_layers.summarize_tensors(norm_tensors, tag="norms/")
# Expand dimension 1 and 2
return tf.expand_dims(decoder_output, axis=1)
def mlp(feature, hparams, name="mlp"):
"""Multi layer perceptron with dropout and relu activation."""
with tf.variable_scope(name, "mlp", values=[feature]):
num_mlp_layers = hparams.num_mlp_layers
mlp_size = hparams.mlp_size
for _ in range(num_mlp_layers):
feature = common_layers.dense(feature, mlp_size, activation=None)
utils.collect_named_outputs("norms", "mlp_feature",
tf.norm(feature, axis=-1))
feature = common_layers.layer_norm(feature)
feature = tf.nn.relu(feature)
feature = tf.nn.dropout(feature, keep_prob=1. - hparams.dropout)
return feature
def mlp(feature, hparams, name="mlp"):
"""Multi layer perceptron with dropout and relu activation."""
with tf.variable_scope(name, "mlp", values=[feature]):
num_mlp_layers = hparams.num_mlp_layers
mlp_size = hparams.mlp_size
for _ in range(num_mlp_layers):
feature = common_layers.dense(feature, mlp_size, activation=None)
utils.collect_named_outputs("norms", "mlp_feature",
tf.norm(feature, axis=-1))
feature = common_layers.layer_norm(feature)
feature = tf.nn.relu(feature)
feature = tf.nn.dropout(feature, keep_prob=1.-hparams.dropout)
return feature
def _compute_edge_transforms(node_states,
depth,
num_transforms,
name="transform"):
"""Helper function that computes transformation for keys and values.
Let B be the number of batches.
Let N be the number of nodes in the graph.
Let D be the size of the node hidden states.
Let K be the size of the attention keys/queries (total_key_depth).
Let V be the size of the attention values (total_value_depth).
Let T be the total number of transforms (num_transforms).
Computes the transforms for keys or values for attention.
* For each node N_j and edge type t, a key K_jt of size K is computed. When an
edge of type t goes from node N_j to any other node, K_jt is the key that is
in the attention process.
* For each node N_j and edge type t, a value V_jt of size V is computed. When
an edge of type t goes from node N_j to node N_i, Attention(Q_i, K_jt)
produces a weight w_ijt. The message sent along this edge is w_ijt * V_jt.
Args:
node_states: A tensor of shape [B, L, D]
depth: An integer (K or V)
num_transforms: An integer (T),
name: A name for the function
Returns:
x: A The attention keys or values for each node and edge type
(shape [B, N*T, K or V])
"""
node_shapes = common_layers.shape_list(node_states)
x = common_layers.dense(
node_states,
depth * num_transforms,
use_bias=False,
name=name)
batch = node_shapes[0] # B.
length = node_shapes[1] # N.
# Making the fourth dimension explicit by separating the vectors of size
# K*T (in k) and V*T (in v) into two-dimensional matrices with shape [K, T]
# (in k) and [V, T] in v.
#
x = tf.reshape(x, [batch, length, num_transforms, depth])
# Flatten out the fourth dimension.
x = tf.reshape(x, [batch, length * num_transforms, depth])
return x
def _compute_edge_transforms(node_states,
depth,
num_transforms,
name="transform"):
"""Helper function that computes transformation for keys and values.
Let B be the number of batches.
Let N be the number of nodes in the graph.
Let D be the size of the node hidden states.
Let K be the size of the attention keys/queries (total_key_depth).
Let V be the size of the attention values (total_value_depth).
Let T be the total number of transforms (num_transforms).
Computes the transforms for keys or values for attention.
* For each node N_j and edge type t, a key K_jt of size K is computed. When an
edge of type t goes from node N_j to any other node, K_jt is the key that is
in the attention process.
* For each node N_j and edge type t, a value V_jt of size V is computed. When
an edge of type t goes from node N_j to node N_i, Attention(Q_i, K_jt)
produces a weight w_ijt. The message sent along this edge is w_ijt * V_jt.
Args:
node_states: A tensor of shape [B, L, D]
depth: An integer (K or V)
num_transforms: An integer (T),
name: A name for the function
Returns:
x: A The attention keys or values for each node and edge type
(shape [B, N*T, K or V])
"""
node_shapes = common_layers.shape_list(node_states)
x = common_layers.dense(
node_states,
depth * num_transforms,
use_bias=False,
name=name)
batch = node_shapes[0] # B.
length = node_shapes[1] # N.
# Making the fourth dimension explicit by separating the vectors of size
# K*T (in k) and V*T (in v) into two-dimensional matrices with shape [K, T]
# (in k) and [V, T] in v.
#
x = tf.reshape(x, [batch, length, num_transforms, depth])
# Flatten out the fourth dimension.
x = tf.reshape(x, [batch, length * num_transforms, depth])
return x
def body(self, features):
hp = self.hparams
# pylint: disable=eval-used
if hp.image_input_type == "image":
image_feat = vqa_layers.image_embedding(
features["inputs"],
model_fn=eval(hp.image_model_fn),
trainable=hp.train_resnet,
is_training=hp.mode == tf.estimator.ModeKeys.TRAIN)
else:
image_feat = features["inputs"]
image_feat = common_layers.flatten4d3d(image_feat)
image_model_d = hp.model_d
image_feat = common_layers.dense(image_feat, image_model_d)
utils.collect_named_outputs("norms", "image_feat_after_proj",
tf.norm(image_feat, axis=-1))
question = common_layers.flatten4d3d(features["question"])
utils.collect_named_outputs("norms", "question_embedding",
tf.norm(question, axis=-1))
(encoder_input, encoder_self_attention_bias,
encoder_decoder_attention_bias) = prepare_image_question_encoder(
image_feat, question, hp)
encoder_input = tf.nn.dropout(encoder_input,
keep_prob=1. -
hp.layer_prepostprocess_dropout)
encoder_output = image_question_encoder(encoder_input,
encoder_self_attention_bias,
hp)
utils.collect_named_outputs("norms", "encoder_output",
tf.norm(encoder_output, axis=-1))
# scale query by sqrt(model_d)
query = tf.get_variable("query", [hp.model_d]) * hp.model_d**0.5
query = tf.expand_dims(tf.expand_dims(query, axis=0), axis=0)
batch_size = common_layers.shape_list(encoder_input)[0]
query = tf.tile(query, [batch_size, 1, 1])
query = tf.nn.dropout(query,
keep_prob=1. - hp.layer_prepostprocess_dropout)
decoder_output = decoder(query, encoder_output, None,
encoder_decoder_attention_bias, hp)
utils.collect_named_outputs("norms", "decoder_output",
tf.norm(decoder_output, axis=-1))
norm_tensors = utils.convert_collection_to_dict("norms")
vqa_layers.summarize_tensors(norm_tensors, tag="norms/")
# Expand dimension 1 and 2
return tf.expand_dims(decoder_output, axis=1)
def BilinearAttention(c, q, c_mask, filters, name, norm=False):
q = tf.expand_dims(dense(q,
filters,
name=name,
reuse=tf.AUTO_REUSE,
kernel_initializer=initializer(),
kernel_regularizer=regularizer),
axis=-1) # [bs, dim, 1]
if norm:
q = layer_norm(q)
cq = tf.squeeze(tf.matmul(c, q),
axis=-1) # [bs, c_len, dim] * [bs, dim, 1] -> [bs, c_len]
cq = exp_mask(cq, c_mask)
return cq
def universal_transformer_all_steps_so_far(layer_inputs, step, hparams,
ffn_unit, attention_unit):
"""universal_transformer.
It uses an attention mechanism-flipped vertically-
over all the states from previous steps to generate the new_state.
Args:
layer_inputs:
- state: state
- memory: contains states from all the previous steps.
step: indicating number of steps take so far
hparams: model hyper-parameters.
ffn_unit: feed-forward unit
attention_unit: multi-head attention unit
Returns:
layer_output:
new_state: new state
memory: contains states from all the previous steps.
"""
_, inputs, memory = layer_inputs
all_states = memory
# get the states up to the current step (non-zero part of the memory)
states_so_far = all_states[:step, :, :, :]
states_so_far_weights = tf.nn.softmax(common_layers.dense(
states_so_far, (hparams.hidden_size if hparams.dwa_elements else 1),
activation=None,
use_bias=True),
axis=-1)
# # get summary of the step weights
# step_weightes = tf.unstack(states_so_far_weights, axis=0, name="step_weightes")
# for step_i, step_w in enumerate(step_weightes):
# tf.contrib.summary.scalar("step_%d_weight:"%step_i,
# tf.reduce_mean(step_w))
# prepare the state as the summary of
state_to_be_transformed = tf.reduce_sum(
(states_so_far * states_so_far_weights), axis=0)
state_to_be_transformed = universal_transformer_util.step_preprocess(
state_to_be_transformed, step, hparams)
new_state = ffn_unit(attention_unit(state_to_be_transformed))
# add the new state to the memory
memory = universal_transformer_util.fill_memory_slot(
memory, new_state, step + 1)
return new_state, inputs, memory
def body(self, features):
hp = self.hparams
model_fn = resnet_v1_152
if hp.image_model_fn != "resnet_v1_152":
model_fn = eval(hp.image_model_fn) # pylint: disable=eval-used
if hp.image_input_type == "image":
image_feat = vqa_layers.image_embedding(
features["inputs"],
model_fn=model_fn,
trainable=hp.train_resnet,
is_training=hp.mode == tf.estimator.ModeKeys.TRAIN)
else:
image_feat = features["inputs"]
if hp.image_feat_size:
image_feat = common_layers.dense(image_feat, hp.image_feat_size)
# apply layer normalization and dropout on image_feature
utils.collect_named_outputs("norms", "image_feat_before_l2",
tf.norm(image_feat, axis=-1))
image_feat = common_layers.l2_norm(image_feat)
utils.collect_named_outputs("norms", "image_feat_after_l2",
tf.norm(image_feat, axis=-1))
image_feat = tf.nn.dropout(image_feat, keep_prob=1.-hp.dropout)
query = question_encoder(features["question"], hp)
utils.collect_named_outputs("norms", "query",
tf.norm(query, axis=-1))
image_ave = attn(image_feat, query, hp)
utils.collect_named_outputs("norms", "image_ave",
tf.norm(image_ave, axis=-1))
image_question = tf.concat([image_ave, query], axis=1)
utils.collect_named_outputs("norms", "image_question",
tf.norm(image_question, axis=-1))
image_question = tf.nn.dropout(image_question, 1. - hp.dropout)
output = mlp(image_question, hp)
utils.collect_named_outputs("norms", "output",
tf.norm(output, axis=-1))
norm_tensors = utils.convert_collection_to_dict("norms")
vqa_layers.summarize_tensors(norm_tensors, tag="norms/")
# Expand dimension 1 and 2
return tf.expand_dims(tf.expand_dims(output, axis=1), axis=2)
def body(self, features):
hp = self.hparams
# pylint: disable=eval-used
if hp.image_input_type == "image":
image_feat = vqa_layers.image_embedding(
features["inputs"],
model_fn=eval(hp.image_model_fn),
trainable=hp.train_resnet,
is_training=hp.mode == tf.estimator.ModeKeys.TRAIN)
else:
image_feat = features["inputs"]
if hp.image_feat_size:
image_feat = common_layers.dense(image_feat, hp.image_feat_size)
# apply layer normalization and dropout on image_feature
utils.collect_named_outputs("norms", "image_feat_before_l2",
tf.norm(image_feat, axis=-1))
image_feat = common_layers.l2_norm(image_feat)
utils.collect_named_outputs("norms", "image_feat_after_l2",
tf.norm(image_feat, axis=-1))
image_feat = tf.nn.dropout(image_feat, keep_prob=1.-hp.dropout)
query = question_encoder(features["question"], hp)
utils.collect_named_outputs("norms", "query",
tf.norm(query, axis=-1))
image_ave = attn(image_feat, query, hp)
utils.collect_named_outputs("norms", "image_ave",
tf.norm(image_ave, axis=-1))
image_question = tf.concat([image_ave, query], axis=1)
utils.collect_named_outputs("norms", "image_question",
tf.norm(image_question, axis=-1))
image_question = tf.nn.dropout(image_question, 1. - hp.dropout)
output = mlp(image_question, hp)
utils.collect_named_outputs("norms", "output",
tf.norm(output, axis=-1))
norm_tensors = utils.convert_collection_to_dict("norms")
vqa_layers.summarize_tensors(norm_tensors, tag="norms/")
# Expand dimension 1 and 2
return tf.expand_dims(tf.expand_dims(output, axis=1), axis=2)
def body(self, features):
target_seq = features["targets_raw"]
target_seq = tf.identity(target_seq, "target_seq")
hp = self._hparams
actions_seq, grid_output_top_list = self.forward_path(features)
# argmaxed_grid_output = tf.identity(tf.argmax(grid_output_top_list, axis=-1), "argmaxed_grid_output")
# Match the output to the target modality (shape, dimension)
output_vocab_size = self._problem_hparams.vocab_size["targets"]
logits = common_layers.dense(grid_output_top_list, output_vocab_size)
# Compute loss
self.compute_loss(target_seq, logits)
# Flip the logits for evaluation during training
final = infer_flipped_outputs(logits, hp.list_size)
return {"targets": final}, self._additional_loss
def _compute_edge_transforms(node_states,
depth,
num_transforms,
ignore_zero=True,
name="transform"):
"""Helper function that computes transformation for keys and values.
Let B be the number of batches.
Let N be the number of nodes in the graph.
Let D be the size of the node hidden states.
Let K be the size of the attention keys/queries (total_key_depth).
Let V be the size of the attention values (total_value_depth).
Let T be the total number of transforms (num_transforms).
Computes the transforms for keys or values for attention.
* For each node N_j and edge type t, a key K_jt of size K is computed. When an
edge of type t goes from node N_j to any other node, K_jt is the key that is
in the attention process.
* For each node N_j and edge type t, a value V_jt of size V is computed. When
an edge of type t goes from node N_j to node N_i, Attention(Q_i, K_jt)
produces a weight w_ijt. The message sent along this edge is w_ijt * V_jt.
Args:
node_states: A tensor of shape [B, L, D]
depth: An integer (K or V)
num_transforms: An integer (T),
ignore_zero: A boolean to ignore 0 edge
name: A name for the function
Returns:
x: A The attention keys or values for each node and edge type
(shape [B, N*T, K or V])
"""
node_shapes = common_layers.shape_list(node_states)
nonignored_transforms = num_transforms - int(ignore_zero)
x = common_layers.dense(node_states,
depth * nonignored_transforms,
use_bias=False,
name=name)
batch = node_shapes[0] # B.
length = node_shapes[1] # N.
# Making the fourth dimension explicit by separating the vectors of size
# K*T (in k) and V*T (in v) into two-dimensional matrices with shape [K, T]
# (in k) and [V, T] in v.
#
# This reshape is only necessary when ignore_zero is True (for the padding
# step that follows).
x = tf.reshape(x, [batch, length, nonignored_transforms, depth])
# If we previously ignored edge type 0, then we need to pad the keys and
# values to take this additional edge type into account. To do so, we
# pad the third dimension of k and v (which has size T-1 if ignore_zero is
# True) to size T with zeroes.
if ignore_zero:
x = tf.pad(x, [[0, 0], [0, 0], [1, 0], [0, 0]])
# Flatten out the fourth dimension.
x = tf.reshape(x, [batch, length * num_transforms, depth])
return x
def sparse_message_pass(node_states,
adjacency_matrices,
num_edge_types,
hidden_size,
use_bias=True,
average_aggregation=False,
name="sparse_ggnn"):
"""One message-passing step for a GNN with a sparse adjacency matrix.
Implements equation 2 (the message passing step) in
[Li et al. 2015](https://arxiv.org/abs/1511.05493).
N = The number of nodes in each batch.
H = The size of the hidden states.
T = The number of edge types.
Args:
node_states: Initial states of each node in the graph. Shape is [N, H].
adjacency_matrices: Adjacency matrix of directed edges for each edge
type. Shape is [N, N, T] (sparse tensor).
num_edge_types: The number of edge types. T.
hidden_size: The size of the hidden state. H.
use_bias: Whether to use bias in the hidden layer.
average_aggregation: How to aggregate the incoming node messages. If
average_aggregation is true, the messages are averaged. If it is false,
they are summed.
name: (optional) The scope within which tf variables should be created.
Returns:
The result of one step of Gated Graph Neural Network (GGNN) message passing.
Shape: [N, H]
"""
n = tf.shape(node_states)[0]
t = num_edge_types
incoming_edges_per_type = tf.sparse_reduce_sum(adjacency_matrices, axis=1)
# Convert the adjacency matrix into shape [T, N, N] - one [N, N] adjacency
# matrix for each edge type. Since sparse tensor multiplication only supports
# two-dimensional tensors, we actually convert the adjacency matrix into a
# [T * N, N] tensor.
adjacency_matrices = tf.sparse_transpose(adjacency_matrices, [2, 0, 1])
adjacency_matrices = tf.sparse_reshape(adjacency_matrices, [t * n, n])
# Multiply the adjacency matrix by the node states, producing a [T * N, H]
# tensor. For each (edge type, node) pair, this tensor stores the sum of
# the hidden states of the node's neighbors over incoming edges of that type.
messages = tf.sparse_tensor_dense_matmul(adjacency_matrices, node_states)
# Rearrange this tensor to have shape [N, T * H]. The incoming states of each
# nodes neighbors are summed by edge type and then concatenated together into
# a single T * H vector.
messages = tf.reshape(messages, [t, n, hidden_size])
messages = tf.transpose(messages, [1, 0, 2])
messages = tf.reshape(messages, [n, t * hidden_size])
# Run each of those T * H vectors through a linear layer that produces
# a vector of size H. This process is equivalent to running each H-sized
# vector through a separate linear layer for each edge type and then adding
# the results together.
#
# Note that, earlier on, we added together all of the states of neighbors
# that were connected by edges of the same edge type. Since addition and
# multiplying by a linear layer are commutative, this process was equivalent
# to running each incoming edge through a linear layer separately and then
# adding everything at the end.
with tf.variable_scope(name, default_name="sparse_ggnn"):
final_node_states = common_layers.dense(messages,
hidden_size,
use_bias=False)
# Multiply the bias by for each edge type by the number of incoming nodes
# of that edge type.
if use_bias:
bias = tf.get_variable("bias",
initializer=tf.zeros([t, hidden_size]))
final_node_states += tf.matmul(incoming_edges_per_type, bias)
if average_aggregation:
incoming_edges = tf.reduce_sum(incoming_edges_per_type,
-1,
keepdims=True)
incoming_edges = tf.tile(incoming_edges, [1, hidden_size])
final_node_states /= incoming_edges + 1e-7
return tf.reshape(final_node_states, [n, hidden_size])
def multihead_mpnn_attention(node_states,
total_key_depth,
total_value_depth,
output_depth,
num_heads,
adjacency_matrix=None,
num_edge_types=5,
num_transforms=None,
use_weighted_sum=False,
name="mpnn_attention"):
"""Multihead scaled-dot-product attention with input/output transformations.
Let B be the number of batches.
Let N be the number of nodes in the graph.
Let D be the size of the node hidden states.
Let K be the size of the attention keys/queries (total_key_depth).
Let V be the size of the attention values (total_value_depth).
Let O be the size of the attention output (output_depth).
Let H be the number of heads (num_heads).
Let T be the total number of transforms (num_transforms).
The key and value depths are split across all of the heads. For example, if
the key depth is 6 and there are three heads, then the key for each head has
depth 2.
Args:
node_states: A Tensor with shape [B, N, D]
total_key_depth: An integer (K).
total_value_depth: An integer (V).
output_depth: An integer (O).
num_heads: An integer (H).
adjacency_matrix: An Tensor of ints with shape [B, T, N, N]. If there is an
edge from node j to node i in batch b, then adjacency_matrix[b, i, j]
contains the type of that edge as an integer. Otherwise, it contains 0.
num_edge_types: An integer indicating number of edge types.
num_transforms: An integer indicating number of transforms (T). If None,
then num_transforms will be equal to num_edge_types.
use_weighted_sum: If False, will only use a single transform per edge type.
Otherwise, use a learned weighted sum of transforms per edge type.
name: A string.
Returns:
The result of the attention transformation. The output shape is [B, N, O].
Raises:
ValueError: if the key depth or value depth are not divisible by the
number of attention heads.
"""
if total_key_depth % num_heads != 0:
raise ValueError("Key depth (%d) must be divisible by the number of "
"attention heads (%d)." %
(total_key_depth, num_heads))
if total_value_depth % num_heads != 0:
raise ValueError("Value depth (%d) must be divisible by the number of "
"attention heads (%d)." %
(total_value_depth, num_heads))
with tf.variable_scope(name,
default_name="multihead_mpnn_attention",
values=[node_states]):
# If not explicitly set, use num_transforms set to num_edge_types.
num_transforms = (num_edge_types
if num_transforms is None else num_transforms)
# Create the query for each node's incoming edges.
# Create the keys/values for each node for each possible outgoing edge type.
q, k, v = compute_mpnn_qkv(node_states, total_key_depth,
total_value_depth, num_transforms)
q_shape = tf.shape(q) # As above, q_shape is [B, N, K].
# Divides each query/key/value into separate heads. Specifically, the
# query/key/value for each (batch, node) pair (i.e., the third dimensions
# of q, k, and v) are broken into H separate pieces. These pieces are used
# as the separate attention heads. The resulting tensors have shape
# [B, H, N, ?/H], where ? = K, K*T or V*T as appropriate.
q = common_attention.split_heads(q, num_heads) # Shape [B, H, N, K/H].
k = common_attention.split_heads(k,
num_heads) # Shape [B, H, N, K*T/H].
v = common_attention.split_heads(v,
num_heads) # Shape [B, H, N, V*T/H].
key_depth_per_head = total_key_depth // num_heads
# Ensures that the logits don't have too large of a magnitude.
q *= key_depth_per_head**-0.5
# Rearrange the dimensions so that the head is first. This will make
# subsequent steps easier (we loop over the head).
q = tf.transpose(q, [1, 0, 2, 3]) # Shape [H, B, N, K/H].
k = tf.transpose(k, [1, 0, 2, 3]) # Shape [H, B, N, K*T/H].
v = tf.transpose(v, [1, 0, 2, 3]) # Shape [H, B, N, V*T/H].
# Split the keys and values into separate per-edge-type keys and values.
k = tf.reshape(k, [
num_heads, q_shape[0], q_shape[1], num_transforms,
total_key_depth // num_heads
]) # Shape [H, B, N, T, K/H].
k = tf.transpose(k, [0, 1, 3, 2, 4]) # Shape [H, B, T, N, K/H].
v = tf.reshape(v, [
num_heads, q_shape[0], q_shape[1], num_transforms,
total_value_depth // num_heads
]) # Shape [H, B, N, T, V/H].
v = tf.transpose(v, [0, 1, 3, 2, 4]) # Shape [H, B, T, N, V/H].
# Perform attention for each head and combine the results into a list.
# head_outputs stores a list of tensors, each with shape [1, B, N, V/H].
# The last dimension contains the values computed for each attention head.
# Each value was determined by computing attention over all of the
# incoming edges for node n, weighting the incoming values accordingly,
# and adding those weighted values together.
head_outputs = []
for head_id in range(num_heads):
output = dot_product_mpnn_attention(
q[head_id],
k[head_id],
v[head_id],
adjacency_matrix,
num_edge_types,
num_transforms=num_transforms,
use_weighted_sum=use_weighted_sum)
# Store this result in the list of attention results for each head.
# The call to expand_dims gives output shape [1, B, N, V/H], which will
# come in handy when we combine the heads together.
head_outputs.append(tf.expand_dims(output, axis=0))
# Combine the heads together into one tensor and rearrange the dimensions.
x = tf.concat(head_outputs, axis=0) # Shape [H, B, N, V/H].
x = tf.transpose(x, [1, 0, 2, 3]) # Shape [B, H, N, V/H].
# Concatenate the values produced by each head together into one vector.
x = common_attention.combine_heads(x) # Shape [B, N, V].
# A fully-connected linear layer to convert from the value vectors of size V
# to output vectors of length O (the appropriate output length).
x = common_layers.dense(x,
output_depth,
use_bias=False,
name="output_transform")
return x
def multihead_graph_attention(query_antecedent,
memory_antecedent,
bias,
total_key_depth,
total_value_depth,
output_depth,
num_heads,
dropout_rate,
image_shapes=None,
attention_type="edge_vector",
name="multihead_graph_attention",
save_weights_to=None,
make_image_summary=True,
dropout_broadcast_dims=None,
adjacency_matrix=None,
num_edge_types=5,
vars_3d=False,
**kwargs):
"""Multihead scaled-dot-product attention with input/output transformations.
Args:
query_antecedent: a Tensor with shape [batch, length_q, channels]
memory_antecedent: a Tensor with shape [batch, length_m, channels] or None
bias: bias Tensor (see attention_bias())
total_key_depth: an integer
total_value_depth: an integer
output_depth: an integer
num_heads: an integer dividing total_key_depth and total_value_depth
dropout_rate: a floating point number
image_shapes: optional tuple of integer scalars.
see comments for attention_image_summary()
attention_type: a string, either "dot_product", "dot_product_relative",
"local_mask_right", "local_unmasked", "masked_dilated_1d",
"unmasked_dilated_1d", graph, or any attention function
with the signature (query, key, value, **kwargs)
name: an optional string.
save_weights_to: an optional dictionary to capture attention weights
for vizualization; the weights tensor will be appended there under
a string key created from the variable scope (including name).
make_image_summary: Whether to make an attention image summary.
dropout_broadcast_dims: an optional list of integers less than 4
specifying in which dimensions to broadcast the dropout decisions.
saves memory.
adjacency_matrix: an optional tensor of shape [batch, len_q, len_q]
containing edge vectors for attention
num_edge_types: number of edge types, an int
vars_3d: use 3-dimensional variables for input/output transformations
**kwargs (dict): Parameters for the attention function
Returns:
The result of the attention transformation. The output shape is
[batch_size, length_q, output_depth]
Raises:
ValueError: if the key depth or value depth are not divisible by the
number of attention heads.
"""
if total_key_depth % num_heads != 0:
raise ValueError("Key depth (%d) must be divisible by the number of "
"attention heads (%d)." %
(total_key_depth, num_heads))
if total_value_depth % num_heads != 0:
raise ValueError("Value depth (%d) must be divisible by the number of "
"attention heads (%d)." %
(total_value_depth, num_heads))
vars_3d_num_heads = num_heads if vars_3d else None
with tf.variable_scope(name,
default_name="multihead_attention",
values=[query_antecedent, memory_antecedent]):
q, k, v = common_attention.compute_qkv(
query_antecedent,
memory_antecedent,
total_key_depth,
total_value_depth,
vars_3d_num_heads=vars_3d_num_heads)
q = common_attention.split_heads(q, num_heads)
k = common_attention.split_heads(k, num_heads)
v = common_attention.split_heads(v, num_heads)
key_depth_per_head = total_key_depth // num_heads
if not vars_3d:
q *= key_depth_per_head**-0.5
additional_returned_value = None
if callable(
attention_type): # Generic way to extend multihead_attention
x = attention_type(q, k, v, **kwargs)
if isinstance(x, tuple):
x, additional_returned_value = x # Unpack
elif attention_type == "edge_vector":
x = graph_attention(q,
k,
v,
bias,
dropout_rate,
image_shapes,
save_weights_to=save_weights_to,
make_image_summary=make_image_summary,
dropout_broadcast_dims=dropout_broadcast_dims,
adjacency_matrix=adjacency_matrix,
num_edge_types=num_edge_types)
x = common_attention.combine_heads(x)
# Set last dim specifically.
x.set_shape(x.shape.as_list()[:-1] + [total_value_depth])
if vars_3d:
o_var = tf.get_variable(
"o", [num_heads, total_value_depth // num_heads, output_depth])
o_var = tf.reshape(o_var, [total_value_depth, output_depth])
x = tf.tensordot(x, o_var, axes=1)
else:
x = common_layers.dense(x,
output_depth,
use_bias=False,
name="output_transform")
if additional_returned_value is not None:
return x, additional_returned_value
return x
def compute_mpnn_qkv(node_states, total_key_depth, total_value_depth,
num_transforms):
"""Computes query, key and value for edge matrices.
Let B be the number of batches.
Let N be the number of nodes in the graph.
Let D be the size of the node hidden states.
Let K be the size of the attention keys/queries (total_key_depth).
Let V be the size of the attention values (total_value_depth).
Let T be the total number of transforms (num_transforms).
Computes the queries, keys, and values for attention.
* For each node N_i in the graph, a query Q_i of size K is computed. This
query is used to determine the relative weights to give to each of the
node's incoming edges.
* For each node N_j and edge type t, a key K_jt of size K is computed. When an
edge of type t goes from node N_j to any other node, K_jt is the key that is
in the attention process.
* For each node N_j and edge type t, a value V_jt of size V is computed. When
an edge of type t goes from node N_j to node N_i, Attention(Q_i, K_jt)
produces a weight w_ijt. The message sent along this edge is w_ijt * V_jt.
Args:
node_states: A Tensor with shape [B, N, D].
total_key_depth: an integer (K).
total_value_depth: an integer (V).
num_transforms: a integer specifying number of transforms (T). This is
typically the number of edge types.
Returns:
q: The attention queries for each destination node (shape [B, N, K]).
k: The attention keys for each node and edge type (shape [B, N*T, K]).
v: The attention values for each node and edge type (shape [B, N*T, V]).
"""
# node_states is initially a tensor with shape [B, N, D]. The call to dense
# creates a D x K kernel that serves as a fully-connected layer.
#
# For each possible batch b and node n in the first two dimensions of
# node_states, the corresponding size-D vector (the third dimension of
# node_states) is the hidden state for node n in batch b. Each of these size-D
# vectors is multiplied by the kernel to produce an attention query of size K.
# The result is a tensor of size [B, N, K] containing the attention queries
# for each node in each batch.
q = common_layers.dense(node_states,
total_key_depth,
use_bias=False,
name="q_mpnn")
# Creates the attention keys in a manner similar to the process of creating
# the attention queries. One key is created for each type of outgoing edge the
# corresponding node might have, meaning k will have shape [B, N, K*T].
k = _compute_edge_transforms(node_states,
total_key_depth,
num_transforms,
name="k_mpnn")
v = _compute_edge_transforms(node_states,
total_value_depth,
num_transforms,
name="v_mpnn")
return q, k, v
def sparse_message_pass(node_states,
adjacency_matrices,
num_edge_types,
hidden_size,
use_bias=True,
average_aggregation=False,
name="sparse_ggnn"):
"""One message-passing step for a GNN with a sparse adjacency matrix.
Implements equation 2 (the message passing step) in
[Li et al. 2015](https://arxiv.org/abs/1511.05493).
N = The number of nodes in each batch.
H = The size of the hidden states.
T = The number of edge types.
Args:
node_states: Initial states of each node in the graph. Shape is [N, H].
adjacency_matrices: Adjacency matrix of directed edges for each edge
type. Shape is [N, N, T] (sparse tensor).
num_edge_types: The number of edge types. T.
hidden_size: The size of the hidden state. H.
use_bias: Whether to use bias in the hidden layer.
average_aggregation: How to aggregate the incoming node messages. If
average_aggregation is true, the messages are averaged. If it is false,
they are summed.
name: (optional) The scope within which tf variables should be created.
Returns:
The result of one step of Gated Graph Neural Network (GGNN) message passing.
Shape: [N, H]
"""
n = tf.shape(node_states)[0]
t = num_edge_types
incoming_edges_per_type = tf.sparse_reduce_sum(adjacency_matrices, axis=1)
# Convert the adjacency matrix into shape [T, N, N] - one [N, N] adjacency
# matrix for each edge type. Since sparse tensor multiplication only supports
# two-dimensional tensors, we actually convert the adjacency matrix into a
# [T * N, N] tensor.
adjacency_matrices = tf.sparse_transpose(adjacency_matrices, [2, 0, 1])
adjacency_matrices = tf.sparse_reshape(adjacency_matrices, [t * n, n])
# Multiply the adjacency matrix by the node states, producing a [T * N, H]
# tensor. For each (edge type, node) pair, this tensor stores the sum of
# the hidden states of the node's neighbors over incoming edges of that type.
messages = tf.sparse_tensor_dense_matmul(adjacency_matrices, node_states)
# Rearrange this tensor to have shape [N, T * H]. The incoming states of each
# nodes neighbors are summed by edge type and then concatenated together into
# a single T * H vector.
messages = tf.reshape(messages, [t, n, hidden_size])
messages = tf.transpose(messages, [1, 0, 2])
messages = tf.reshape(messages, [n, t * hidden_size])
# Run each of those T * H vectors through a linear layer that produces
# a vector of size H. This process is equivalent to running each H-sized
# vector through a separate linear layer for each edge type and then adding
# the results together.
#
# Note that, earlier on, we added together all of the states of neighbors
# that were connected by edges of the same edge type. Since addition and
# multiplying by a linear layer are commutative, this process was equivalent
# to running each incoming edge through a linear layer separately and then
# adding everything at the end.
with tf.variable_scope(name, default_name="sparse_ggnn"):
final_node_states = common_layers.dense(
messages, hidden_size, use_bias=False)
# Multiply the bias by for each edge type by the number of incoming nodes
# of that edge type.
if use_bias:
bias = tf.get_variable("bias", initializer=tf.zeros([t, hidden_size]))
final_node_states += tf.matmul(incoming_edges_per_type, bias)
if average_aggregation:
incoming_edges = tf.reduce_sum(incoming_edges_per_type, -1, keepdims=True)
incoming_edges = tf.tile(incoming_edges, [1, hidden_size])
final_node_states /= incoming_edges + 1e-7
return final_node_states
def body(self, features):
hp = self.hparams
# pylint: disable=eval-used
if hp.image_input_type == "image":
image_feat = vqa_layers.image_embedding(
features["inputs"],
model_fn=eval(hp.image_model_fn),
trainable=hp.train_resnet,
is_training=hp.mode == tf.estimator.ModeKeys.TRAIN)
else:
image_feat = features["inputs"]
image_feat = common_layers.flatten4d3d(image_feat)
image_hidden_size = hp.image_hidden_size or hp.hidden_size
if hp.image_feat_preprocess_proj:
image_feat = common_layers.dense(image_feat, image_hidden_size)
utils.collect_named_outputs("norms", "image_feat_after_proj",
tf.norm(image_feat, axis=-1))
else:
assert image_hidden_size == 2048
image_feat = tf.nn.dropout(image_feat,
keep_prob=1. -
hp.layer_prepostprocess_dropout)
if hp.image_feat_encode:
image_feat = image_encoder(image_feat, hp)
utils.collect_named_outputs("norms", "image_feat_encoded",
tf.norm(image_feat, axis=-1))
else:
image_feat = common_layers.layer_norm(image_feat)
utils.collect_named_outputs("norms", "image_feat_after_layer",
tf.norm(image_feat, axis=-1))
question = common_layers.flatten4d3d(features["question"])
utils.collect_named_outputs("norms", "question_embedding",
tf.norm(question, axis=-1))
question, question_self_attention_bias = prepare_question_encoder(
question, hp)
question = tf.nn.dropout(question,
keep_prob=1. -
hp.layer_prepostprocess_dropout)
query = question_encoder(question, question_self_attention_bias, hp)
utils.collect_named_outputs("norms", "query_encode",
tf.norm(query, axis=-1))
query = (query + tf.expand_dims(
tf.squeeze(question_self_attention_bias, [1, 2]), axis=2))
query = tf.reduce_max(query, axis=1)
utils.collect_named_outputs("norms", "query_maxpool",
tf.norm(query, axis=-1))
# query = common_layers.l2_norm(query)
# utils.collect_named_outputs("norms", "query_after_l2",
# tf.norm(query, axis=-1))
image_ave = attn(image_feat, query, hp)
utils.collect_named_outputs("norms", "image_ave",
tf.norm(image_ave, axis=-1))
if hp.multimodal_combine == "concat":
image_question = tf.concat([image_ave, query], axis=1)
elif hp.multimodal_combine == "sum":
image_question = image_ave + query
elif hp.multimodal_combine == "product":
image_question = image_ave * query
utils.collect_named_outputs("norms", "image_question",
tf.norm(image_question, axis=-1))
image_question = tf.nn.dropout(image_question, 1. - hp.dropout)
output = mlp(image_question, hp)
utils.collect_named_outputs("norms", "output", tf.norm(output,
axis=-1))
norm_tensors = utils.convert_collection_to_dict("norms")
vqa_layers.summarize_tensors(norm_tensors, tag="norms/")
# Expand dimension 1 and 2
return tf.expand_dims(tf.expand_dims(output, axis=1), axis=2)
def multihead_attention(query_antecedent,
memory_antecedent,
bias,
total_key_depth,
total_value_depth,
output_depth,
num_heads,
dropout_rate,
shared_rel=False,
max_relative_position=None,
image_shapes=None,
attention_type="dot_product",
block_length=128,
block_width=128,
q_filter_width=1,
kv_filter_width=1,
q_padding="VALID",
kv_padding="VALID",
cache=None,
gap_size=0,
num_memory_blocks=2,
name="multihead_attention",
save_weights_to=None,
make_image_summary=True,
dropout_broadcast_dims=None,
max_length=None,
vars_3d=False,
scale_dotproduct=True,
**kwargs):
"""Multihead scaled-dot-product attention with input/output transformations.
Args:
query_antecedent: a Tensor with shape [batch, length_q, channels]
memory_antecedent: a Tensor with shape [batch, length_m, channels] or None
bias: bias Tensor (see attention_bias())
total_key_depth: an integer
total_value_depth: an integer
output_depth: an integer
num_heads: an integer dividing total_key_depth and total_value_depth
dropout_rate: a floating point number
shared_rel: boolean to share relative embeddings
max_relative_position: Maximum distance between inputs to generate
unique relation embeddings for. Only relevant
when using "dot_product_relative" attention.
image_shapes: optional tuple of integer scalars.
see comments for attention_image_summary()
attention_type: a string, either "dot_product", "dot_product_relative",
"local_mask_right", "local_unmasked", "masked_dilated_1d",
"unmasked_dilated_1d", graph, or any attention function
with the signature (query, key, value, **kwargs)
block_length: an integer - relevant for "local_mask_right"
block_width: an integer - relevant for "local_unmasked"
q_filter_width: An integer specifying how wide you want the query to be.
kv_filter_width: An integer specifying how wide you want the keys and values
to be.
q_padding: One of "VALID", "SAME" or "LEFT". Default is VALID: No padding.
kv_padding: One of "VALID", "SAME" or "LEFT". Default is "VALID":
no padding.
cache: dict containing Tensors which are the results of previous
attentions, used for fast decoding. Expects the dict to contrain two
keys ('k' and 'v'), for the initial call the values for these keys
should be empty Tensors of the appropriate shape.
'k' [batch_size, 0, key_channels]
'v' [batch_size, 0, value_channels]
gap_size: Integer option for dilated attention to indicate spacing between
memory blocks.
num_memory_blocks: Integer option to indicate how many memory blocks to look
at.
name: an optional string.
save_weights_to: an optional dictionary to capture attention weights
for vizualization; the weights tensor will be appended there under
a string key created from the variable scope (including name).
make_image_summary: Whether to make an attention image summary.
dropout_broadcast_dims: an optional list of integers less than 4
specifying in which dimensions to broadcast the dropout decisions.
saves memory.
max_length: an integer - needed by relative attention
vars_3d: use 3-dimensional variables for input/output transformations
scale_dotproduct: whether to normalize the attention product.
**kwargs (dict): Parameters for the attention function
Caching:
WARNING: For decoder self-attention, i.e. when memory_antecedent == None,
the caching assumes that the bias contains future masking.
The caching works by saving all the previous key and value values so that
you are able to send just the last query location to this attention
function. I.e. if the cache dict is provided it assumes the query is of the
shape [batch_size, 1, hidden_dim] rather than the full memory.
Returns:
The result of the attention transformation. The output shape is
[batch_size, length_q, hidden_dim]
unless the cache dict is provided in which case only the last memory
position is calculated and the output shape is [batch_size, 1, hidden_dim]
Optionally returns an additional loss parameters (ex: load balance loss for
the experts) returned by the attention_type function.
Raises:
ValueError: if the key depth or value depth are not divisible by the
number of attention heads.
"""
if total_key_depth % num_heads != 0:
raise ValueError("Key depth (%d) must be divisible by the number of "
"attention heads (%d)." % (total_key_depth, num_heads))
if total_value_depth % num_heads != 0:
raise ValueError("Value depth (%d) must be divisible by the number of "
"attention heads (%d)." % (total_value_depth, num_heads))
vars_3d_num_heads = num_heads if vars_3d else 0
with tf.variable_scope(name, default_name="multihead_attention",
values=[query_antecedent, memory_antecedent]):
if cache is None or memory_antecedent is None:
q, k, v = common_attention.compute_qkv(
query_antecedent, memory_antecedent,
total_key_depth, total_value_depth, q_filter_width,
kv_filter_width, q_padding, kv_padding,
vars_3d_num_heads=vars_3d_num_heads)
if cache is not None:
if attention_type != "dot_product":
# TODO(petershaw): Support caching when using relative position
# representations, i.e. "dot_product_relative" attention.
raise NotImplementedError(
"Caching is not guaranteed to work with attention types other than"
" dot_product.")
if bias is None:
raise ValueError("Bias required for caching. See function docstring "
"for details.")
if memory_antecedent is not None:
# Encoder-Decoder Attention Cache
q = common_attention.compute_attention_component(
query_antecedent, total_key_depth,
q_filter_width, q_padding, "q",
vars_3d_num_heads=vars_3d_num_heads)
k = cache["k_encdec"]
v = cache["v_encdec"]
else:
k = common_attention.split_heads(k, num_heads)
v = common_attention.split_heads(v, num_heads)
decode_loop_step = kwargs.get("decode_loop_step")
if decode_loop_step is None:
k = cache["k"] = tf.concat([cache["k"], k], axis=2)
v = cache["v"] = tf.concat([cache["v"], v], axis=2)
else:
# Inplace update is required for inference on TPU.
# Inplace_ops only supports inplace_update on the first dimension.
# The performance of current implementation is better than updating
# the tensor by adding the result of matmul(one_hot,
# update_in_current_step)
tmp_k = tf.transpose(cache["k"], perm=[2, 0, 1, 3])
tmp_k = inplace_ops.alias_inplace_update(
tmp_k, decode_loop_step, tf.squeeze(k, axis=2))
k = cache["k"] = tf.transpose(tmp_k, perm=[1, 2, 0, 3])
tmp_v = tf.transpose(cache["v"], perm=[2, 0, 1, 3])
tmp_v = inplace_ops.alias_inplace_update(
tmp_v, decode_loop_step, tf.squeeze(v, axis=2))
v = cache["v"] = tf.transpose(tmp_v, perm=[1, 2, 0, 3])
q = common_attention.split_heads(q, num_heads)
if cache is None:
k = common_attention.split_heads(k, num_heads)
v = common_attention.split_heads(v, num_heads)
key_depth_per_head = total_key_depth // num_heads
if not vars_3d:
if scale_dotproduct:
q *= key_depth_per_head**-0.5
additional_returned_value = None
if callable(attention_type): # Generic way to extend multihead_attention
x = attention_type(q, k, v, **kwargs)
if isinstance(x, tuple):
x, additional_returned_value = x # Unpack
elif attention_type == "dot_product":
x = common_attention.dot_product_attention(
q, k, v, bias, dropout_rate, image_shapes,
save_weights_to=save_weights_to,
make_image_summary=make_image_summary,
dropout_broadcast_dims=dropout_broadcast_dims)
elif attention_type == "dot_product_relative":
x = common_attention.dot_product_attention_relative(
q,
k,
v,
bias,
max_relative_position,
dropout_rate,
image_shapes,
make_image_summary=make_image_summary)
elif attention_type == "dot_product_relative_v2":
x = common_attention.dot_product_self_attention_relative_v2(
q,
k,
v,
bias,
max_length,
dropout_rate,
image_shapes,
make_image_summary=make_image_summary,
dropout_broadcast_dims=dropout_broadcast_dims)
elif attention_type == "local_within_block_mask_right":
x = common_attention.masked_within_block_local_attention_1d(
q, k, v, block_length=block_length)
elif attention_type == "rel_local_mask_right":
x = common_attention.masked_rel_local_attention_1d(
q, k, v, block_length=block_length,
make_image_summary=make_image_summary,
dropout_rate=dropout_rate,
share_rel_embed=shared_rel)
elif attention_type == "local_mask_right":
x = common_attention.masked_local_attention_1d(
q,
k,
v,
block_length=block_length,
make_image_summary=make_image_summary)
elif attention_type == "local_unmasked":
x = common_attention.local_attention_1d(
q, k, v, block_length=block_length, filter_width=block_width)
elif attention_type == "masked_dilated_1d":
x = common_attention.masked_dilated_self_attention_1d(
q, k, v, block_length, block_width,
gap_size, num_memory_blocks)
else:
assert attention_type == "unmasked_dilated_1d"
x = common_attention.dilated_self_attention_1d(
q, k, v, block_length, block_width,
gap_size, num_memory_blocks)
x = common_attention.combine_heads(x)
# Set last dim specifically.
x.set_shape(x.shape.as_list()[:-1] + [total_value_depth])
if vars_3d:
o_var = tf.get_variable(
"o", [num_heads, total_value_depth // num_heads, output_depth])
o_var = tf.cast(o_var, x.dtype)
o_var = tf.reshape(o_var, [total_value_depth, output_depth])
x = tf.tensordot(x, o_var, axes=1)
else:
x = common_layers.dense(
x, output_depth, use_bias=False, name="output_transform")
if additional_returned_value is not None:
return x, additional_returned_value
return x
def multihead_attention(query_antecedent,
memory_antecedent,
bias,
total_key_depth,
total_value_depth,
output_depth,
num_heads,
dropout_rate,
shared_rel=False,
max_relative_position=None,
image_shapes=None,
attention_type="dot_product",
block_length=128,
block_width=128,
q_filter_width=1,
kv_filter_width=1,
q_padding="VALID",
kv_padding="VALID",
cache=None,
gap_size=0,
num_memory_blocks=2,
name="multihead_attention",
save_weights_to=None,
make_image_summary=True,
dropout_broadcast_dims=None,
max_length=None,
vars_3d=False,
scale_dotproduct=True,
**kwargs):
"""Multihead scaled-dot-product attention with input/output transformations.
Args:
query_antecedent: a Tensor with shape [batch, length_q, channels]
memory_antecedent: a Tensor with shape [batch, length_m, channels] or None
bias: bias Tensor (see attention_bias())
total_key_depth: an integer
total_value_depth: an integer
output_depth: an integer
num_heads: an integer dividing total_key_depth and total_value_depth
dropout_rate: a floating point number
shared_rel: boolean to share relative embeddings
max_relative_position: Maximum distance between inputs to generate
unique relation embeddings for. Only relevant
when using "dot_product_relative" attention.
image_shapes: optional tuple of integer scalars.
see comments for attention_image_summary()
attention_type: a string, either "dot_product", "dot_product_relative",
"local_mask_right", "local_unmasked", "masked_dilated_1d",
"unmasked_dilated_1d", graph, or any attention function
with the signature (query, key, value, **kwargs)
block_length: an integer - relevant for "local_mask_right"
block_width: an integer - relevant for "local_unmasked"
q_filter_width: An integer specifying how wide you want the query to be.
kv_filter_width: An integer specifying how wide you want the keys and values
to be.
q_padding: One of "VALID", "SAME" or "LEFT". Default is VALID: No padding.
kv_padding: One of "VALID", "SAME" or "LEFT". Default is "VALID":
no padding.
cache: dict containing Tensors which are the results of previous
attentions, used for fast decoding. Expects the dict to contrain two
keys ('k' and 'v'), for the initial call the values for these keys
should be empty Tensors of the appropriate shape.
'k' [batch_size, 0, key_channels]
'v' [batch_size, 0, value_channels]
gap_size: Integer option for dilated attention to indicate spacing between
memory blocks.
num_memory_blocks: Integer option to indicate how many memory blocks to look
at.
name: an optional string.
save_weights_to: an optional dictionary to capture attention weights
for vizualization; the weights tensor will be appended there under
a string key created from the variable scope (including name).
make_image_summary: Whether to make an attention image summary.
dropout_broadcast_dims: an optional list of integers less than 4
specifying in which dimensions to broadcast the dropout decisions.
saves memory.
max_length: an integer - needed by relative attention
vars_3d: use 3-dimensional variables for input/output transformations
scale_dotproduct: whether to normalize the attention product.
**kwargs (dict): Parameters for the attention function
Caching:
WARNING: For decoder self-attention, i.e. when memory_antecedent == None,
the caching assumes that the bias contains future masking.
The caching works by saving all the previous key and value values so that
you are able to send just the last query location to this attention
function. I.e. if the cache dict is provided it assumes the query is of the
shape [batch_size, 1, hidden_dim] rather than the full memory.
Returns:
The result of the attention transformation. The output shape is
[batch_size, length_q, hidden_dim]
unless the cache dict is provided in which case only the last memory
position is calculated and the output shape is [batch_size, 1, hidden_dim]
Optionally returns an additional loss parameters (ex: load balance loss for
the experts) returned by the attention_type function.
Raises:
ValueError: if the key depth or value depth are not divisible by the
number of attention heads.
"""
if total_key_depth % num_heads != 0:
raise ValueError("Key depth (%d) must be divisible by the number of "
"attention heads (%d)." %
(total_key_depth, num_heads))
if total_value_depth % num_heads != 0:
raise ValueError("Value depth (%d) must be divisible by the number of "
"attention heads (%d)." %
(total_value_depth, num_heads))
vars_3d_num_heads = num_heads if vars_3d else 0
with tf.variable_scope(name,
default_name="multihead_attention",
values=[query_antecedent, memory_antecedent]):
if cache is None or memory_antecedent is None:
q, k, v = common_attention.compute_qkv(
query_antecedent,
memory_antecedent,
total_key_depth,
total_value_depth,
q_filter_width,
kv_filter_width,
q_padding,
kv_padding,
vars_3d_num_heads=vars_3d_num_heads)
if cache is not None:
if attention_type != "dot_product":
# TODO(petershaw): Support caching when using relative position
# representations, i.e. "dot_product_relative" attention.
raise NotImplementedError(
"Caching is not guaranteed to work with attention types other than"
" dot_product.")
if bias is None:
raise ValueError(
"Bias required for caching. See function docstring "
"for details.")
if memory_antecedent is not None:
# Encoder-Decoder Attention Cache
q = common_attention.compute_attention_component(
query_antecedent,
total_key_depth,
q_filter_width,
q_padding,
"q",
vars_3d_num_heads=vars_3d_num_heads)
k = cache["k_encdec"]
v = cache["v_encdec"]
else:
k = common_attention.split_heads(k, num_heads)
v = common_attention.split_heads(v, num_heads)
decode_loop_step = kwargs.get("decode_loop_step")
if decode_loop_step is None:
k = cache["k"] = tf.concat([cache["k"], k], axis=2)
v = cache["v"] = tf.concat([cache["v"], v], axis=2)
else:
# Inplace update is required for inference on TPU.
# Inplace_ops only supports inplace_update on the first dimension.
# The performance of current implementation is better than updating
# the tensor by adding the result of matmul(one_hot,
# update_in_current_step)
tmp_k = tf.transpose(cache["k"], perm=[2, 0, 1, 3])
tmp_k = inplace_ops.alias_inplace_update(
tmp_k, decode_loop_step, tf.squeeze(k, axis=2))
k = cache["k"] = tf.transpose(tmp_k, perm=[1, 2, 0, 3])
tmp_v = tf.transpose(cache["v"], perm=[2, 0, 1, 3])
tmp_v = inplace_ops.alias_inplace_update(
tmp_v, decode_loop_step, tf.squeeze(v, axis=2))
v = cache["v"] = tf.transpose(tmp_v, perm=[1, 2, 0, 3])
q = common_attention.split_heads(q, num_heads)
if cache is None:
k = common_attention.split_heads(k, num_heads)
v = common_attention.split_heads(v, num_heads)
key_depth_per_head = total_key_depth // num_heads
if not vars_3d:
if scale_dotproduct:
q *= key_depth_per_head**-0.5
additional_returned_value = None
if callable(
attention_type): # Generic way to extend multihead_attention
x = attention_type(q, k, v, **kwargs)
if isinstance(x, tuple):
x, additional_returned_value = x # Unpack
elif attention_type == "dot_product":
x = common_attention.dot_product_attention(
q,
k,
v,
bias,
dropout_rate,
image_shapes,
save_weights_to=save_weights_to,
make_image_summary=make_image_summary,
dropout_broadcast_dims=dropout_broadcast_dims)
elif attention_type == "dot_product_relative":
x = common_attention.dot_product_attention_relative(
q,
k,
v,
bias,
max_relative_position,
dropout_rate,
image_shapes,
make_image_summary=make_image_summary)
elif attention_type == "dot_product_relative_v2":
x = common_attention.dot_product_self_attention_relative_v2(
q,
k,
v,
bias,
max_length,
dropout_rate,
image_shapes,
make_image_summary=make_image_summary,
dropout_broadcast_dims=dropout_broadcast_dims)
elif attention_type == "local_within_block_mask_right":
x = common_attention.masked_within_block_local_attention_1d(
q, k, v, block_length=block_length)
elif attention_type == "rel_local_mask_right":
x = common_attention.masked_rel_local_attention_1d(
q,
k,
v,
block_length=block_length,
make_image_summary=make_image_summary,
dropout_rate=dropout_rate,
share_rel_embed=shared_rel)
elif attention_type == "local_mask_right":
x = common_attention.masked_local_attention_1d(
q,
k,
v,
block_length=block_length,
make_image_summary=make_image_summary)
elif attention_type == "local_unmasked":
x = common_attention.local_attention_1d(q,
k,
v,
block_length=block_length,
filter_width=block_width)
elif attention_type == "masked_dilated_1d":
x = common_attention.masked_dilated_self_attention_1d(
q, k, v, block_length, block_width, gap_size,
num_memory_blocks)
else:
assert attention_type == "unmasked_dilated_1d"
x = common_attention.dilated_self_attention_1d(
q, k, v, block_length, block_width, gap_size,
num_memory_blocks)
x = common_attention.combine_heads(x)
# Set last dim specifically.
x.set_shape(x.shape.as_list()[:-1] + [total_value_depth])
if vars_3d:
o_var = tf.get_variable(
"o", [num_heads, total_value_depth // num_heads, output_depth])
o_var = tf.cast(o_var, x.dtype)
o_var = tf.reshape(o_var, [total_value_depth, output_depth])
x = tf.tensordot(x, o_var, axes=1)
else:
x = common_layers.dense(x,
output_depth,
use_bias=False,
name="output_transform")
if additional_returned_value is not None:
return x, additional_returned_value
return x
def multihead_mpnn_attention(node_states,
total_key_depth,
total_value_depth,
output_depth,
num_heads,
adjacency_matrix=None,
num_edge_types=5,
num_transforms=None,
use_weighted_sum=False,
name="mpnn_attention"):
"""Multihead scaled-dot-product attention with input/output transformations.
Let B be the number of batches.
Let N be the number of nodes in the graph.
Let D be the size of the node hidden states.
Let K be the size of the attention keys/queries (total_key_depth).
Let V be the size of the attention values (total_value_depth).
Let O be the size of the attention output (output_depth).
Let H be the number of heads (num_heads).
Let T be the total number of transforms (num_transforms).
The key and value depths are split across all of the heads. For example, if
the key depth is 6 and there are three heads, then the key for each head has
depth 2.
Args:
node_states: A Tensor with shape [B, N, D]
total_key_depth: An integer (K).
total_value_depth: An integer (V).
output_depth: An integer (O).
num_heads: An integer (H).
adjacency_matrix: An Tensor of ints with shape [B, T, N, N]. If there is an
edge from node j to node i in batch b, then adjacency_matrix[b, i, j]
contains the type of that edge as an integer. Otherwise, it contains 0.
num_edge_types: An integer indicating number of edge types.
num_transforms: An integer indicating number of transforms (T). If None,
then num_transforms will be equal to num_edge_types.
use_weighted_sum: If False, will only use a single transform per edge type.
Otherwise, use a learned weighted sum of transforms per edge type.
name: A string.
Returns:
The result of the attention transformation. The output shape is [B, N, O].
Raises:
ValueError: if the key depth or value depth are not divisible by the
number of attention heads.
"""
if total_key_depth % num_heads != 0:
raise ValueError("Key depth (%d) must be divisible by the number of "
"attention heads (%d)." % (total_key_depth, num_heads))
if total_value_depth % num_heads != 0:
raise ValueError("Value depth (%d) must be divisible by the number of "
"attention heads (%d)." % (total_value_depth, num_heads))
with tf.variable_scope(
name, default_name="multihead_mpnn_attention", values=[node_states]):
# If not explicitly set, use num_transforms set to num_edge_types.
num_transforms = (
num_edge_types if num_transforms is None else num_transforms)
# Create the query for each node's incoming edges.
# Create the keys/values for each node for each possible outgoing edge type.
q, k, v = compute_mpnn_qkv(
node_states,
total_key_depth,
total_value_depth,
num_transforms)
q_shape = tf.shape(q) # As above, q_shape is [B, N, K].
# Divides each query/key/value into separate heads. Specifically, the
# query/key/value for each (batch, node) pair (i.e., the third dimensions
# of q, k, and v) are broken into H separate pieces. These pieces are used
# as the separate attention heads. The resulting tensors have shape
# [B, H, N, ?/H], where ? = K, K*T or V*T as appropriate.
q = common_attention.split_heads(q, num_heads) # Shape [B, H, N, K/H].
k = common_attention.split_heads(k, num_heads) # Shape [B, H, N, K*T/H].
v = common_attention.split_heads(v, num_heads) # Shape [B, H, N, V*T/H].
key_depth_per_head = total_key_depth // num_heads
# Ensures that the logits don't have too large of a magnitude.
q *= key_depth_per_head**-0.5
# Rearrange the dimensions so that the head is first. This will make
# subsequent steps easier (we loop over the head).
q = tf.transpose(q, [1, 0, 2, 3]) # Shape [H, B, N, K/H].
k = tf.transpose(k, [1, 0, 2, 3]) # Shape [H, B, N, K*T/H].
v = tf.transpose(v, [1, 0, 2, 3]) # Shape [H, B, N, V*T/H].
# Split the keys and values into separate per-edge-type keys and values.
k = tf.reshape(k, [
num_heads, q_shape[0], q_shape[1], num_transforms,
total_key_depth // num_heads
]) # Shape [H, B, N, T, K/H].
k = tf.transpose(k, [0, 1, 3, 2, 4]) # Shape [H, B, T, N, K/H].
v = tf.reshape(v, [
num_heads, q_shape[0], q_shape[1], num_transforms,
total_value_depth // num_heads
]) # Shape [H, B, N, T, V/H].
v = tf.transpose(v, [0, 1, 3, 2, 4]) # Shape [H, B, T, N, V/H].
# Perform attention for each head and combine the results into a list.
# head_outputs stores a list of tensors, each with shape [1, B, N, V/H].
# The last dimension contains the values computed for each attention head.
# Each value was determined by computing attention over all of the
# incoming edges for node n, weighting the incoming values accordingly,
# and adding those weighted values together.
head_outputs = []
for head_id in range(num_heads):
output = dot_product_mpnn_attention(
q[head_id],
k[head_id],
v[head_id],
adjacency_matrix,
num_edge_types,
num_transforms=num_transforms,
use_weighted_sum=use_weighted_sum)
# Store this result in the list of attention results for each head.
# The call to expand_dims gives output shape [1, B, N, V/H], which will
# come in handy when we combine the heads together.
head_outputs.append(tf.expand_dims(output, axis=0))
# Combine the heads together into one tensor and rearrange the dimensions.
x = tf.concat(head_outputs, axis=0) # Shape [H, B, N, V/H].
x = tf.transpose(x, [1, 0, 2, 3]) # Shape [B, H, N, V/H].
# Concatenate the values produced by each head together into one vector.
x = common_attention.combine_heads(x) # Shape [B, N, V].
# A fully-connected linear layer to convert from the value vectors of size V
# to output vectors of length O (the appropriate output length).
x = common_layers.dense(
x, output_depth, use_bias=False, name="output_transform")
return x
def multihead_graph_attention(query_antecedent,
memory_antecedent,
bias,
total_key_depth,
total_value_depth,
output_depth,
num_heads,
dropout_rate,
image_shapes=None,
attention_type="edge_vector",
name="multihead_graph_attention",
save_weights_to=None,
make_image_summary=True,
dropout_broadcast_dims=None,
adjacency_matrix=None,
num_edge_types=5,
vars_3d=False,
**kwargs):
"""Multihead scaled-dot-product attention with input/output transformations.
Args:
query_antecedent: a Tensor with shape [batch, length_q, channels]
memory_antecedent: a Tensor with shape [batch, length_m, channels] or None
bias: bias Tensor (see attention_bias())
total_key_depth: an integer
total_value_depth: an integer
output_depth: an integer
num_heads: an integer dividing total_key_depth and total_value_depth
dropout_rate: a floating point number
image_shapes: optional tuple of integer scalars.
see comments for attention_image_summary()
attention_type: a string, either "dot_product", "dot_product_relative",
"local_mask_right", "local_unmasked", "masked_dilated_1d",
"unmasked_dilated_1d", graph, or any attention function
with the signature (query, key, value, **kwargs)
name: an optional string.
save_weights_to: an optional dictionary to capture attention weights
for vizualization; the weights tensor will be appended there under
a string key created from the variable scope (including name).
make_image_summary: Whether to make an attention image summary.
dropout_broadcast_dims: an optional list of integers less than 4
specifying in which dimensions to broadcast the dropout decisions.
saves memory.
adjacency_matrix: an optional tensor of shape [batch, len_q, len_q]
containing edge vectors for attention
num_edge_types: number of edge types, an int
vars_3d: use 3-dimensional variables for input/output transformations
**kwargs (dict): Parameters for the attention function
Returns:
The result of the attention transformation. The output shape is
[batch_size, length_q, output_depth]
Raises:
ValueError: if the key depth or value depth are not divisible by the
number of attention heads.
"""
if total_key_depth % num_heads != 0:
raise ValueError("Key depth (%d) must be divisible by the number of "
"attention heads (%d)." % (total_key_depth, num_heads))
if total_value_depth % num_heads != 0:
raise ValueError("Value depth (%d) must be divisible by the number of "
"attention heads (%d)." % (total_value_depth, num_heads))
vars_3d_num_heads = num_heads if vars_3d else None
with tf.variable_scope(
name,
default_name="multihead_attention",
values=[query_antecedent, memory_antecedent]):
q, k, v = common_attention.compute_qkv(
query_antecedent,
memory_antecedent,
total_key_depth,
total_value_depth,
vars_3d_num_heads=vars_3d_num_heads)
q = common_attention.split_heads(q, num_heads)
k = common_attention.split_heads(k, num_heads)
v = common_attention.split_heads(v, num_heads)
key_depth_per_head = total_key_depth // num_heads
if not vars_3d:
q *= key_depth_per_head**-0.5
additional_returned_value = None
if callable(attention_type): # Generic way to extend multihead_attention
x = attention_type(q, k, v, **kwargs)
if isinstance(x, tuple):
x, additional_returned_value = x # Unpack
elif attention_type == "edge_vector":
x = graph_attention(
q,
k,
v,
bias,
dropout_rate,
image_shapes,
save_weights_to=save_weights_to,
make_image_summary=make_image_summary,
dropout_broadcast_dims=dropout_broadcast_dims,
adjacency_matrix=adjacency_matrix,
num_edge_types=num_edge_types)
x = common_attention.combine_heads(x)
# Set last dim specifically.
x.set_shape(x.shape.as_list()[:-1] + [total_value_depth])
if vars_3d:
o_var = tf.get_variable(
"o", [num_heads, total_value_depth // num_heads, output_depth])
o_var = tf.reshape(o_var, [total_value_depth, output_depth])
x = tf.tensordot(x, o_var, axes=1)
else:
x = common_layers.dense(
x, output_depth, use_bias=False, name="output_transform")
if additional_returned_value is not None:
return x, additional_returned_value
return x
def compute_mpnn_qkv(node_states,
total_key_depth,
total_value_depth,
num_edge_types,
ignore_zero=True):
"""Computes query, key and value for edge matrices.
Let B be the number of batches.
Let N be the number of nodes in the graph.
Let D be the size of the node hidden states.
Let K be the size of the attention keys/queries (total_key_depth).
Let V be the size of the attention values (total_value_depth).
Let T be the total number of edge types (num_edge_types).
Computes the queries, keys, and values for attention.
* For each node N_i in the graph, a query Q_i of size K is computed. This
query is used to determine the relative weights to give to each of the
node's incoming edges.
* For each node N_j and edge type t, a key K_jt of size K is computed. When an
edge of type t goes from node N_j to any other node, K_jt is the key that is
in the attention process.
* For each node N_j and edge type t, a value V_jt of size V is computed. When
an edge of type t goes from node N_j to node N_i, Attention(Q_i, K_jt)
produces a weight w_ijt. The message sent along this edge is w_ijt * V_jt.
Args:
node_states: A Tensor with shape [B, N, D].
total_key_depth: an integer (K).
total_value_depth: an integer (V).
num_edge_types: a integer specifying number of edge types (T).
ignore_zero: If true, then edge type 0 will not be considered. Equivalent
to having a linear transformation of all 0's for edge type 0. All queries,
keys, and values for edge type 0 will be all 0's.
Returns:
q: The attention queries for each destination node (shape [B, N, K]).
k: The attention keys for each node and edge type (shape [B, N*T, K]).
v: The attention values for each node and edge type (shape [B, N*T, V]).
"""
# node_states is initially a tensor with shape [B, N, D]. The call to dense
# creates a D x K kernel that serves as a fully-connected layer.
#
# For each possible batch b and node n in the first two dimensions of
# node_states, the corresponding size-D vector (the third dimension of
# node_states) is the hidden state for node n in batch b. Each of these size-D
# vectors is multiplied by the kernel to produce an attention query of size K.
# The result is a tensor of size [B, N, K] containing the attention queries
# for each node in each batch.
q = common_layers.dense(node_states,
total_key_depth,
use_bias=False,
name="q_mpnn")
q_shape = common_layers.shape_list(q) # As above, q_shape = [B, N, K].
# T (or T-1 if ignore_zero).
nonignored_edge_types = num_edge_types - int(ignore_zero)
# Creates the attention keys in a manner similar to the process of creating
# the attention queries. One key is created for each type of outgoing edge the
# corresponding node might have, meaning k will have shape [B, N, K*T].
k = common_layers.dense(node_states,
total_key_depth * nonignored_edge_types,
use_bias=False,
name="k_mpnn")
# The values over which self-attention is performed. They are created in
# a manner largely identical to that of the keys.
v = common_layers.dense(node_states,
total_value_depth * nonignored_edge_types,
use_bias=False,
name="v_mpnn")
batch = q_shape[0] # B.
length = q_shape[1] # N.
# Making the fourth dimension explicit by separating the vectors of size
# K*T (in k) and V*T (in v) into two-dimensional matrices with shape [K, T]
# (in k) and [V, T] in v.
#
# This reshape is only necessary when ignore_zero is True (for the padding
# step that follows).
k = tf.reshape(k, [batch, length, nonignored_edge_types, total_key_depth])
v = tf.reshape(
v, [q_shape[0], q_shape[1], nonignored_edge_types, total_value_depth])
# If we previously ignored edge type 0, then we need to pad the keys and
# values to take this additional edge type into account. To do so, we
# pad the third dimension of k and v (which has size T-1 if ignore_zero is
# True) to size T with zeroes.
if ignore_zero:
k = tf.pad(k, [[0, 0], [0, 0], [1, 0], [0, 0]])
v = tf.pad(v, [[0, 0], [0, 0], [1, 0], [0, 0]])
# Flatten out the fourth dimension.
k = tf.reshape(k,
[q_shape[0], q_shape[1] * num_edge_types, total_key_depth])
v = tf.reshape(
v, [q_shape[0], q_shape[1] * num_edge_types, total_value_depth])
return q, k, v
def compute_attention_component(antecedent,
total_depth,
filter_width=1,
padding="VALID",
name="c",
vars_3d_num_heads=0,
sparsity_technique=None,
threshold=3.0,
training=True,
clip_alpha=None,
initial_sparsity=None,
split_heads=False,
num_heads=None):
"""Computes attention compoenent (query, key or value).
Args:
antecedent: a Tensor with shape [batch, length, channels]
total_depth: an integer
filter_width: An integer specifying how wide you want the attention
component to be.
padding: One of "VALID", "SAME" or "LEFT". Default is VALID: No padding.
name: a string specifying scope name.
vars_3d_num_heads: an optional integer (if we want to use 3d variables)
sparsity_technique: technique used for sparsifying weights.
threshold: log alpha threshold used for evaluation with variational dropout.
training: whether model is being trained or not.
clip_alpha: alpha clipping threshold for variational dropout.
initial_sparsity: initial sparsity level for lottery ticket &
scratch experiments.
split_heads: Whether to prune each head separately.
num_heads: The number of heads in the attention module.
Returns:
c : [batch, length, depth] tensor
"""
# We don't support 3d attention variables or filter_width > 1 with sparsity
# techniques
assert not sparsity_technique or (not vars_3d_num_heads
and filter_width == 1)
if vars_3d_num_heads > 0:
assert filter_width == 1
input_depth = antecedent.get_shape().as_list()[-1]
depth_per_head = total_depth // vars_3d_num_heads
initializer_stddev = input_depth**-0.5
if "q" in name:
initializer_stddev *= depth_per_head**-0.5
var = tf.get_variable(
name,
[input_depth, vars_3d_num_heads, total_depth // vars_3d_num_heads],
initializer=tf.random_normal_initializer(
stddev=initializer_stddev))
var = tf.cast(var, antecedent.dtype)
var = tf.reshape(var, [input_depth, total_depth])
return tf.tensordot(antecedent, var, axes=1)
if filter_width == 1:
if sparsity_technique:
if split_heads:
# Prune each heads weights separately so that they are free
# to have different weight magnitude distributions.
if num_heads is None:
raise ValueError(
"`num_heads` must be set for split head pruning.")
if total_depth % num_heads != 0:
raise ValueError(
"`total_depth` must be divisible by `num_heads`.")
input_depth = antecedent.get_shape().as_list()[-1]
depth_per_head = int(total_depth / num_heads)
masked_head_weights = []
for head_id in range(num_heads):
head_name = name + "_shard_{}".format(head_id)
with tf.variable_scope(head_name) as vs:
head_weights = tf.get_variable(
"kernel", [input_depth, depth_per_head])
masked_head_weights.append(
pruning.apply_mask(head_weights, vs))
component_weights = tf.concat(masked_head_weights, axis=1)
# compute the full component result
return tf.tensordot(antecedent, component_weights, axes=1)
else:
return common_sparse.dense(
antecedent,
total_depth,
use_bias=False,
sparsity_technique=sparsity_technique,
threshold=threshold,
training=training,
clip_alpha=clip_alpha,
name=name,
initial_sparsity=initial_sparsity)
else:
return common_layers.dense(antecedent,
total_depth,
use_bias=False,
name=name)
else:
return common_layers.conv1d(antecedent,
total_depth,
filter_width,
padding=padding,
name=name)
def multihead_attention(query_antecedent,
memory_antecedent,
bias,
total_key_depth,
total_value_depth,
output_depth,
num_heads,
dropout_rate,
attention_type="dot_product",
image_shapes=None,
q_filter_width=1,
kv_filter_width=1,
q_padding="VALID",
kv_padding="VALID",
cache=None,
name="multihead_attention",
save_weights_to=None,
make_image_summary=True,
dropout_broadcast_dims=None,
vars_3d=False,
sparsity_technique=None,
threshold=3.0,
training=True,
clip_alpha=None,
initial_sparsity=None,
split_heads=False,
**kwargs):
"""Multihead scaled-dot-product attention with input/output transformations.
Args:
query_antecedent: a Tensor with shape [batch, length_q, channels]
memory_antecedent: a Tensor with shape [batch, length_m, channels] or None
bias: bias Tensor (see attention_bias())
total_key_depth: an integer
total_value_depth: an integer
output_depth: an integer
num_heads: an integer dividing total_key_depth and total_value_depth
dropout_rate: a floating point number
attention_type: a string, either "dot_product", "dot_product_relative",
"local_mask_right", "local_unmasked", "masked_dilated_1d",
"unmasked_dilated_1d", graph, or any attention function
with the signature (query, key, value, **kwargs)
image_shapes: optional tuple of integer scalars.
see comments for attention_image_summary()
q_filter_width: An integer specifying how wide you want the query to be.
kv_filter_width: An integer specifying how wide you want the keys and values
to be.
q_padding: One of "VALID", "SAME" or "LEFT". Default is VALID: No padding.
kv_padding: One of "VALID", "SAME" or "LEFT". Default is "VALID":
no padding.
cache: dict containing Tensors which are the results of previous
attentions, used for fast decoding. Expects the dict to contrain two
keys ('k' and 'v'), for the initial call the values for these keys
should be empty Tensors of the appropriate shape.
'k' [batch_size, 0, key_channels]
'v' [batch_size, 0, value_channels]
name: an optional string.
save_weights_to: an optional dictionary to capture attention weights
for vizualization; the weights tensor will be appended there under
a string key created from the variable scope (including name).
make_image_summary: Whether to make an attention image summary.
dropout_broadcast_dims: an optional list of integers less than 4
specifying in which dimensions to broadcast the dropout decisions.
saves memory.
vars_3d: use 3-dimensional variables for input/output transformations
sparsity_technique: technique used for sparsifying weights.
threshold: log alpha threshold used for evaluation with variational dropout.
training: whether model is being trained or not.
clip_alpha: alpha clipping threshold for variational dropout.
initial_sparsity: initial sparsity level for lottery ticket &
scratch experiments.
split_heads: Whether to prune each head separately.
**kwargs (dict): Parameters for the attention function
Caching:
WARNING: For decoder self-attention, i.e. when memory_antecedent == None,
the caching assumes that the bias contains future masking.
The caching works by saving all the previous key and value values so that
you are able to send just the last query location to this attention
function. I.e. if the cache dict is provided it assumes the query is of the
shape [batch_size, 1, hidden_dim] rather than the full memory.
Returns:
The result of the attention transformation. The output shape is
[batch_size, length_q, hidden_dim]
unless the cache dict is provided in which case only the last memory
position is calculated and the output shape is [batch_size, 1, hidden_dim]
Optionally returns an additional loss parameters (ex: load balance loss for
the experts) returned by the attention_type function.
Raises:
ValueError: if the key depth or value depth are not divisible by the
number of attention heads.
"""
if total_key_depth % num_heads != 0:
raise ValueError("Key depth (%d) must be divisible by the number of "
"attention heads (%d)." %
(total_key_depth, num_heads))
if total_value_depth % num_heads != 0:
raise ValueError("Value depth (%d) must be divisible by the number of "
"attention heads (%d)." %
(total_value_depth, num_heads))
if vars_3d:
raise ValueError("3d attention variables not supported.")
if attention_type != "dot_product":
raise ValueError(
"Sparse multihead attention only supports dot_product attention.")
vars_3d_num_heads = 0
with tf.variable_scope(name,
default_name="multihead_attention",
values=[query_antecedent, memory_antecedent]):
if cache is None or memory_antecedent is None:
q, k, v = compute_qkv(query_antecedent,
memory_antecedent,
total_key_depth,
total_value_depth,
q_filter_width,
kv_filter_width,
q_padding,
kv_padding,
vars_3d_num_heads=vars_3d_num_heads,
sparsity_technique=sparsity_technique,
threshold=threshold,
training=training,
clip_alpha=clip_alpha,
initial_sparsity=initial_sparsity,
split_heads=split_heads,
num_heads=num_heads)
if cache is not None:
if bias is None:
raise ValueError(
"Bias required for caching. See function docstring "
"for details.")
if memory_antecedent is not None:
# Encoder-Decoder Attention Cache
q = compute_attention_component(
query_antecedent,
total_key_depth,
q_filter_width,
q_padding,
"q",
vars_3d_num_heads=vars_3d_num_heads,
sparsity_technique=sparsity_technique,
threshold=threshold,
training=training,
clip_alpha=clip_alpha,
initial_sparsity=initial_sparsity,
split_heads=split_heads,
num_heads=num_heads)
k = cache["k_encdec"]
v = cache["v_encdec"]
else:
k = common_attention.split_heads(k, num_heads)
v = common_attention.split_heads(v, num_heads)
decode_loop_step = kwargs.get("decode_loop_step")
if decode_loop_step is None:
k = cache["k"] = tf.concat([cache["k"], k], axis=2)
v = cache["v"] = tf.concat([cache["v"], v], axis=2)
else:
# Inplace update is required for inference on TPU.
# Inplace_ops only supports inplace_update on the first dimension.
# The performance of current implementation is better than updating
# the tensor by adding the result of matmul(one_hot,
# update_in_current_step)
tmp_k = tf.transpose(cache["k"], perm=[2, 0, 1, 3])
tmp_k = inplace_ops.alias_inplace_update(
tmp_k, decode_loop_step, tf.squeeze(k, axis=2))
k = cache["k"] = tf.transpose(tmp_k, perm=[1, 2, 0, 3])
tmp_v = tf.transpose(cache["v"], perm=[2, 0, 1, 3])
tmp_v = inplace_ops.alias_inplace_update(
tmp_v, decode_loop_step, tf.squeeze(v, axis=2))
v = cache["v"] = tf.transpose(tmp_v, perm=[1, 2, 0, 3])
q = common_attention.split_heads(q, num_heads)
if cache is None:
k = common_attention.split_heads(k, num_heads)
v = common_attention.split_heads(v, num_heads)
key_depth_per_head = total_key_depth // num_heads
if not vars_3d:
q *= key_depth_per_head**-0.5
# compute the attention
x = common_attention.dot_product_attention(
q,
k,
v,
bias,
dropout_rate,
image_shapes,
save_weights_to=save_weights_to,
make_image_summary=make_image_summary,
dropout_broadcast_dims=dropout_broadcast_dims)
x = common_attention.combine_heads(x)
# Set last dim specifically.
x.set_shape(x.shape.as_list()[:-1] + [total_value_depth])
if sparsity_technique:
x = common_sparse.dense(x,
output_depth,
use_bias=False,
sparsity_technique=sparsity_technique,
threshold=threshold,
training=training,
clip_alpha=clip_alpha,
name="output_transform",
initial_sparsity=initial_sparsity)
else:
x = common_layers.dense(x,
output_depth,
use_bias=False,
name="output_transform")
return x
def multihead_mpnn_attention(node_states,
total_key_depth,
total_value_depth,
output_depth,
num_heads,
adjacency_matrix=None,
num_edge_types=5,
ignore_zero=True,
name="mpnn_attention"):
"""Multihead scaled-dot-product attention with input/output transformations.
Args:
node_states: A tensor of shape [batch, length, depth]
total_key_depth: An integer for key dimension
total_value_depth: An integer for value dimensions
output_depth: An intger for output dimemsions
num_heads: An integer
adjacency_matrix: An tensor of ints of shape [batch, length, length]
num_edge_types: An integer indicating number of edge bins
ignore_zero: A flag that says that edge type 0 should be ignored
name: A string
Returns:
The result of the attention transformation. The output shape is
[batch_size, length_q, output_depth]
unless the cache dict is provided in which case only the last memory
position is calculated and the output shape is [batch_size, 1, hidden_dim]
Optionaly returns an additional loss parameters (ex: load balance loss for
the experts) returned by the attention_type function.
Raises:
ValueError: if the key depth or value depth are not divisible by the
number of attention heads.
"""
if total_key_depth % num_heads != 0:
raise ValueError("Key depth (%d) must be divisible by the number of "
"attention heads (%d)." %
(total_key_depth, num_heads))
if total_value_depth % num_heads != 0:
raise ValueError("Value depth (%d) must be divisible by the number of "
"attention heads (%d)." %
(total_value_depth, num_heads))
with tf.variable_scope(name,
default_name="multihead_mpnn_attention",
values=[node_states]):
q, k, v = compute_mpnn_qkv(node_states,
total_key_depth,
total_value_depth,
num_edge_types,
ignore_zero=ignore_zero)
# reshaping k and v for head splitting
q_shape = tf.shape(q)
q = common_attention.split_heads(q, num_heads)
k = common_attention.split_heads(k, num_heads)
v = common_attention.split_heads(v, num_heads)
key_depth_per_head = total_key_depth // num_heads
q *= key_depth_per_head**-0.5
# make the heads dimension leading. We will loop over heads.
q = tf.transpose(q, [1, 0, 2, 3])
k = tf.transpose(k, [1, 0, 2, 3])
v = tf.transpose(v, [1, 0, 2, 3])
# putting edge as the dimension after batch for k and v
# k and v will be [heads, batch, num_edge_types, length, depth]
k = tf.reshape(k, [
num_heads, q_shape[0], q_shape[1], num_edge_types,
total_key_depth // num_heads
])
k = tf.transpose(k, [0, 1, 3, 2, 4])
v = tf.reshape(v, [
num_heads, q_shape[0], q_shape[1], num_edge_types,
total_value_depth // num_heads
])
v = tf.transpose(v, [0, 1, 3, 2, 4])
# doing attention separately for each head
head_outputs = []
for head_id in range(num_heads):
output = dot_product_mpnn_attention(q[head_id], k[head_id],
v[head_id], adjacency_matrix,
num_edge_types)
head_outputs.append(tf.expand_dims(output, axis=0))
# making x = [heads, batch, length, total_value_depth//num_heads]
x = tf.concat(head_outputs, axis=0)
x = tf.transpose(x, [1, 0, 2, 3])
# making x [batch, length, depth]
x = common_attention.combine_heads(x)
x = common_layers.dense(x,
output_depth,
use_bias=False,
name="output_transform")
return x
def _compute(inp, depth, filter_width, padding, name):
if filter_width == 1:
return common_layers.dense(inp, depth, use_bias=False, name=name)
else:
return common_layers.conv1d(inp, depth, filter_width, padding, name=name)
def hierarchical_attention_network_encoder(
encoder_input,
encoder_self_attention_bias,
contexts,
context_self_attention_biases,
features,
hparams,
name="hierarchical_attention_network_encoder",
save_weights_to=None,
make_image_summary=True,
losses=None):
input_x = encoder_input
context_xs = {}
for context_name in contexts:
context_xs[context_name] = contexts[context_name]
context_paddings = {}
context_nonpaddings = {}
context_pad_removers = {}
attention_dropout_broadcast_dims = (
common_layers.comma_separated_string_to_integer_list(
getattr(hparams, "attention_dropout_broadcast_dims", "")))
with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
input_padding = common_attention.attention_bias_to_padding(
encoder_self_attention_bias)
input_nonpadding = 1.0 - input_padding
for context_name in context_self_attention_biases:
context_paddings[
context_name] = common_attention.attention_bias_to_padding(
context_self_attention_biases[context_name])
context_nonpaddings[
context_name] = 1.0 - context_paddings[context_name]
input_pad_remover = None
for context_name in context_paddings:
context_pad_removers[context_name] = None
if hparams.use_pad_remover and not common_layers.is_xla_compiled():
input_pad_remover = expert_utils.PadRemover(input_padding)
for context_name in context_paddings:
context_pad_removers[context_name] = expert_utils.PadRemover(
context_paddings[context_name])
temp_hparam = tf.contrib.training.HParams(
) # copy hparams except num_hidden_layers -> num_hidden_layers - 1
for key, val in hparams.values().items():
temp_hparam.add_hparam(key, val)
temp_hparam.set_hparam("num_hidden_layers",
hparams.num_hidden_layers - 1)
encoder_output = transformer_with_contexts_layers.transformer_encoder(
input_x,
encoder_self_attention_bias,
temp_hparam,
nonpadding=features_to_nonpadding(features, "inputs"),
save_weights_to=save_weights_to,
make_image_summary=make_image_summary)
context_encoded_outputs = {}
for context_name in context_xs:
context_encoded_outputs[
context_name] = transformer_with_contexts_layers.transformer_encoder(
context_xs[context_name],
context_self_attention_biases[context_name],
hparams,
nonpadding=features_to_nonpadding(features, context_name),
save_weights_to=save_weights_to,
make_image_summary=make_image_summary)
with tf.variable_scope('word_abstraction', reuse=tf.AUTO_REUSE):
encoder_word_level_query = common_layers.dense(
encoder_output, hparams.hidden_size) # q_w = f_w(h_t)
encoder_word_level_abstraction = {}
for context_name in context_encoded_outputs:
encoder_word_level_abstraction[
context_name] = transformer_with_contexts_layers.multihead_attention(
common_layers.layer_preprocess(
encoder_word_level_query, hparams),
context_encoded_outputs[context_name],
context_self_attention_biases[context_name],
hparams.attention_key_channels or hparams.hidden_size,
hparams.attention_value_channels
or hparams.hidden_size,
hparams.hidden_size,
hparams.num_heads,
hparams.attention_dropout,
attention_type=hparams.self_attention_type,
save_weights_to=save_weights_to,
make_image_summary=make_image_summary,
max_relative_position=hparams.max_relative_position,
dropout_broadcast_dims=attention_dropout_broadcast_dims,
max_length=hparams.get("max_length"),
vars_3d=hparams.get("attention_variables_3d")) # s^j,
sentence_information = tf.concat([
encoder_word_level_abstraction[context_name]
for context_name in encoder_word_level_abstraction
],
axis=1)
with tf.variable_scope('sentence_abstraction', reuse=tf.AUTO_REUSE):
encoder_sentence_level_query = common_layers.dense(
encoder_output, hparams.hidden_size) # q_s = f_s(h_t)
context_padding = common_attention.embedding_to_padding(
sentence_information)
ignore_padding = common_attention.attention_bias_ignore_padding(
context_padding)
contextual_information = transformer_with_contexts_layers.multihead_attention(
common_layers.layer_preprocess(encoder_sentence_level_query,
hparams),
sentence_information,
ignore_padding,
hparams.attention_key_channels or hparams.hidden_size,
hparams.attention_value_channels or hparams.hidden_size,
hparams.hidden_size,
hparams.num_heads,
hparams.attention_dropout,
attention_type=hparams.self_attention_type,
save_weights_to=save_weights_to,
make_image_summary=make_image_summary,
max_relative_position=hparams.max_relative_position,
dropout_broadcast_dims=attention_dropout_broadcast_dims,
max_length=hparams.get("max_length"),
vars_3d=hparams.get("attention_variables_3d")
) # MultiHead(q_s, s^j), [batch, encoder_length, hidden_dim]
contextual_information = common_layers.dense_relu_dense(
contextual_information, hparams.filter_size,
hparams.hidden_size)
with tf.variable_scope('context_gating', reuse=tf.AUTO_REUSE):
gate_lambda = tf.nn.sigmoid(
common_layers.dense(contextual_information,
hparams.hidden_size) +
common_layers.dense(encoder_output, hparams.hidden_size))
encoder_output = gate_lambda * encoder_output + (
1 - gate_lambda) * contextual_information
return common_layers.layer_preprocess(encoder_output, hparams)
def compute_mpnn_qkv(node_states,
total_key_depth,
total_value_depth,
num_transforms):
"""Computes query, key and value for edge matrices.
Let B be the number of batches.
Let N be the number of nodes in the graph.
Let D be the size of the node hidden states.
Let K be the size of the attention keys/queries (total_key_depth).
Let V be the size of the attention values (total_value_depth).
Let T be the total number of transforms (num_transforms).
Computes the queries, keys, and values for attention.
* For each node N_i in the graph, a query Q_i of size K is computed. This
query is used to determine the relative weights to give to each of the
node's incoming edges.
* For each node N_j and edge type t, a key K_jt of size K is computed. When an
edge of type t goes from node N_j to any other node, K_jt is the key that is
in the attention process.
* For each node N_j and edge type t, a value V_jt of size V is computed. When
an edge of type t goes from node N_j to node N_i, Attention(Q_i, K_jt)
produces a weight w_ijt. The message sent along this edge is w_ijt * V_jt.
Args:
node_states: A Tensor with shape [B, N, D].
total_key_depth: an integer (K).
total_value_depth: an integer (V).
num_transforms: a integer specifying number of transforms (T). This is
typically the number of edge types.
Returns:
q: The attention queries for each destination node (shape [B, N, K]).
k: The attention keys for each node and edge type (shape [B, N*T, K]).
v: The attention values for each node and edge type (shape [B, N*T, V]).
"""
# node_states is initially a tensor with shape [B, N, D]. The call to dense
# creates a D x K kernel that serves as a fully-connected layer.
#
# For each possible batch b and node n in the first two dimensions of
# node_states, the corresponding size-D vector (the third dimension of
# node_states) is the hidden state for node n in batch b. Each of these size-D
# vectors is multiplied by the kernel to produce an attention query of size K.
# The result is a tensor of size [B, N, K] containing the attention queries
# for each node in each batch.
q = common_layers.dense(
node_states, total_key_depth, use_bias=False, name="q_mpnn")
# Creates the attention keys in a manner similar to the process of creating
# the attention queries. One key is created for each type of outgoing edge the
# corresponding node might have, meaning k will have shape [B, N, K*T].
k = _compute_edge_transforms(node_states,
total_key_depth,
num_transforms,
name="k_mpnn")
v = _compute_edge_transforms(node_states,
total_value_depth,
num_transforms,
name="v_mpnn")
return q, k, v
def body(self, features):
hp = self.hparams
# pylint: disable=eval-used
if hp.image_input_type == "image":
image_feat = vqa_layers.image_embedding(
features["inputs"],
model_fn=eval(hp.image_model_fn),
trainable=hp.train_resnet,
is_training=hp.mode == tf.estimator.ModeKeys.TRAIN)
else:
image_feat = features["inputs"]
image_feat = common_layers.flatten4d3d(image_feat)
image_hidden_size = hp.image_hidden_size or hp.hidden_size
if hp.image_feat_preprocess_proj:
image_feat = common_layers.dense(image_feat, image_hidden_size)
utils.collect_named_outputs("norms", "image_feat_after_proj",
tf.norm(image_feat, axis=-1))
else:
assert image_hidden_size == 2048
image_feat = tf.nn.dropout(
image_feat, keep_prob=1.-hp.layer_prepostprocess_dropout)
if hp.image_feat_encode:
image_feat = image_encoder(image_feat, hp)
utils.collect_named_outputs("norms", "image_feat_encoded",
tf.norm(image_feat, axis=-1))
else:
image_feat = common_layers.layer_norm(image_feat)
utils.collect_named_outputs("norms", "image_feat_after_layer",
tf.norm(image_feat, axis=-1))
question = common_layers.flatten4d3d(features["question"])
utils.collect_named_outputs("norms", "question_embedding",
tf.norm(question, axis=-1))
question, question_self_attention_bias = prepare_question_encoder(
question, hp)
question = tf.nn.dropout(
question, keep_prob=1.-hp.layer_prepostprocess_dropout)
query = question_encoder(question, question_self_attention_bias, hp)
utils.collect_named_outputs(
"norms", "query_encode", tf.norm(query, axis=-1))
query = (query + tf.expand_dims(
tf.squeeze(question_self_attention_bias, [1, 2]), axis=2))
query = tf.reduce_max(query, axis=1)
utils.collect_named_outputs(
"norms", "query_maxpool", tf.norm(query, axis=-1))
# query = common_layers.l2_norm(query)
# utils.collect_named_outputs("norms", "query_after_l2",
# tf.norm(query, axis=-1))
image_ave = attn(image_feat, query, hp)
utils.collect_named_outputs("norms", "image_ave",
tf.norm(image_ave, axis=-1))
if hp.multimodal_combine == "concat":
image_question = tf.concat([image_ave, query], axis=1)
elif hp.multimodal_combine == "sum":
image_question = image_ave + query
elif hp.multimodal_combine == "product":
image_question = image_ave * query
utils.collect_named_outputs("norms", "image_question",
tf.norm(image_question, axis=-1))
image_question = tf.nn.dropout(image_question, 1. - hp.dropout)
output = mlp(image_question, hp)
utils.collect_named_outputs("norms", "output",
tf.norm(output, axis=-1))
norm_tensors = utils.convert_collection_to_dict("norms")
vqa_layers.summarize_tensors(norm_tensors, tag="norms/")
# Expand dimension 1 and 2
return tf.expand_dims(tf.expand_dims(output, axis=1), axis=2)
def _ffn_layer_multi_inputs(inputs_list,
hparams,
ffn_layer_type="dense",
name="ffn",
kernel_initializer=None,
bias_initializer=None,
activation=None,
pad_remover=None,
preprocess=False,
postprocess=False):
"""Implements a Feed-forward layer with multiple inputs, pad-removing, etc.
Args:
inputs_list: list of input tensors
hparams: hyper-parameters
ffn_layer_type: dense / dense_dropconnect/ dense_relu_dense
name: name
kernel_initializer: kernel initializer
bias_initializer: bias initializer
activation: activation function
pad_remover: pad remover
preprocess: if preprocess the input
postprocess: if postprocess the output
Returns:
a tensor
Raises:
ValueError: Unknown ffn_layer type.
"""
# need at least one inputs
num_inputs = len(inputs_list)
assert num_inputs > 0
if preprocess and num_inputs == 1:
inputs_list[0] = common_layers.layer_preprocess(
inputs_list[0], hparams)
if postprocess:
original_inputs = inputs_list[0]
# the output size is the hidden size of the main inputs
main_input = inputs_list[0]
original_shape = common_layers.shape_list(main_input)
assert hparams.hidden_size == common_layers.shape_list(main_input)[-1]
# all the inputs are in the same shape with main inputs
for inputs in inputs_list:
main_input.get_shape().assert_is_compatible_with(inputs.get_shape())
def remove_pads(x):
original_shape = common_layers.shape_list(x)
# Collapse `x` across examples, and remove padding positions.
x = tf.reshape(x, tf.concat([[-1], original_shape[2:]], axis=0))
x = tf.expand_dims(pad_remover.remove(x), axis=0)
return x
if pad_remover:
for i, inputs in enumerate(inputs_list):
inputs_list[i] = remove_pads(inputs)
ffn_inputs = (inputs_list[0] if len(inputs_list) == 1 else tf.concat(
inputs_list, axis=-1))
if ffn_layer_type == "dense":
output = common_layers.dense(ffn_inputs,
hparams.hidden_size,
name=name,
activation=activation,
use_bias=True,
kernel_initializer=kernel_initializer,
bias_initializer=bias_initializer)
elif ffn_layer_type == "dense_dropconnect":
output = common_layers.dense_dropconnect(
ffn_inputs,
hparams.hidden_size,
name=name,
dropconnect_dropout=hparams.dropconnect_dropout,
output_activation=activation)
postprocess = False # no dropout on the output unit
elif ffn_layer_type == "dense_relu_dense":
output = common_layers.dense_relu_dense(
ffn_inputs,
hparams.filter_size,
hparams.hidden_size,
name=name,
dropout=hparams.relu_dropout,
output_activation=activation,
)
else:
raise ValueError("Unknown ffn_layer type: %s" % ffn_layer_type)
if pad_remover:
# Restore `output` to the original shape of `x`, including padding.
output = tf.reshape(pad_remover.restore(tf.squeeze(output, axis=0)),
original_shape)
if postprocess:
if num_inputs == 1:
output = common_layers.layer_postprocess(original_inputs, output,
hparams)
else: # only dropout (no residual)x
hp = copy.copy(hparams)
hp.layer_postprocess_sequence = hp.layer_postprocess_sequence.replace(
"a", "")
output = common_layers.layer_postprocess(original_inputs, output,
hp)
return output
