def testSpmm_Replicated(self, r, m, k, n, sparsity, use_gpu):
# Helpers to set up the matrices.
connector = connectors.Uniform(sparsity, round_to=4)
initializer = initializers.Uniform()
# Numpy matrices for verification.
mask = connector(initializer([m, k]))
mask[mask != 0] = 1.0
lhs_np = np.expand_dims(mask, axis=0) * initializer([r, m, k])
rhs_np = initializer([r, k, n])
# TensorFlow graph.
topology = sparse_matrix.SparseTopology("topology", mask=mask)
lhs = tf.Variable(np.reshape(lhs_np[lhs_np != 0], [r, -1]),
dtype=tf.float32)
rhs = tf.Variable(rhs_np, dtype=tf.float32)
output = ops.replicated_spmm(lhs, topology, rhs)
# Execute the op and compare the results.
with self.test_session(use_gpu=use_gpu) as sess:
sess.run(tf.global_variables_initializer())
out = sess.run(output)
for i in range(r):
expected_out = np.dot(lhs_np[i, :, :], rhs_np[i, :, :])
self.assertAllClose(out[i, :],
expected_out,
atol=1e-03,
rtol=1e-05)
def testSddmm_Replicated(self, r, m, k, n, sparsity, use_gpu):
# Helpers to set up the matrices.
connector = connectors.Uniform(sparsity)
initializer = initializers.Uniform()
# Numpy matrices for verification.
lhs_np = initializer([r, m, k])
rhs_np = initializer([r, n, k])
output_np = connector(np.ones([m, n]))
# TensorFlow graph.
output_topology = sparse_matrix.SparseTopology("output_topology",
mask=output_np)
lhs = tf.Variable(lhs_np, dtype=tf.float32)
rhs = tf.Variable(rhs_np, dtype=tf.float32)
output = ops.replicated_sddmm(lhs,
rhs,
output_topology,
transpose_rhs=True)
# Execute the op and compare the results.
with self.test_session(use_gpu=use_gpu) as sess:
sess.run(tf.global_variables_initializer())
# Run the replicated sddmm.
v, ro, ci = sess.run([
output, output_topology.row_offsets,
output_topology.column_indices
])
for i in range(r):
expected_output = self.dense_to_scipy(
output_np *
np.dot(lhs_np[i, :, :], np.transpose(rhs_np[i, :, :])))
actual_output = self.sparse_to_scipy(
v[i, :], ro, ci, shape=expected_output.shape)
self.assert_sparse_matrix_equal(actual_output,
expected_output,
atol=1e-03,
rtol=1e-05)
def testSparseSoftmax_Replicated(self, r, m, n, sparsity):
# Helpers to set up the matrices.
connector = connectors.Uniform(sparsity)
initializer = initializers.Uniform()
# Numpy matrix for verification.
mask = connector(np.ones([m, n]))
matrix_np = np.expand_dims(mask, axis=0) * initializer([r, m, n])
# TensorFlow graph.
topology = sparse_matrix.SparseTopology("topology", mask=mask)
values = tf.Variable(np.reshape(matrix_np[matrix_np != 0], [r, -1]),
dtype=tf.float32)
output = ops.replicated_sparse_softmax(values, topology)
with self.test_session(use_gpu=True) as sess:
sess.run(tf.global_variables_initializer())
v, ro, ci = sess.run(
[output, topology.row_offsets, topology.column_indices])
# Zero terms should not contribute to the softmax.
matrix_np[matrix_np == 0] = -1e9
def softmax(x):
maxs = np.expand_dims(x.max(axis=1), axis=1)
exps = np.exp(x - maxs)
return exps / np.expand_dims(np.sum(exps, axis=1), axis=1)
for i in range(r):
expected_output = self.dense_to_scipy(
softmax(matrix_np[i, :, :]))
actual_output = self.sparse_to_scipy(v[i, :], ro, ci,
expected_output.shape)
self.assert_sparse_matrix_equal(actual_output,
expected_output,
atol=1e-03,
rtol=1e-05)
def transformer_decoder(decoder_input,
encoder_output,
decoder_self_attention_bias,
encoder_decoder_attention_bias,
hparams,
cache=None,
decode_loop_step=None,
name="decoder",
nonpadding=None,
save_weights_to=None,
make_image_summary=True,
losses=None,
layer_collection=None,
recurrent_memory_by_layer=None,
chunk_number=None):
"""A stack of transformer layers.
Args:
decoder_input: a Tensor
encoder_output: a Tensor
decoder_self_attention_bias: bias Tensor for self-attention (see
common_attention.attention_bias())
encoder_decoder_attention_bias: bias Tensor for encoder-decoder attention
(see common_attention.attention_bias())
hparams: hyperparameters for model
cache: dict, containing tensors which are the results of previous
attentions, used for fast decoding.
decode_loop_step: An integer, step number of the decoding loop. Only used
for inference on TPU.
name: a string
nonpadding: optional Tensor with shape [batch_size, encoder_length]
indicating what positions are not padding. This is used to mask out
padding in convolutional layers. We generally only need this mask for
"packed" datasets, because for ordinary datasets, no padding is ever
followed by nonpadding.
save_weights_to: an optional dictionary to capture attention weights for
visualization; 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.
losses: optional list onto which to append extra training losses
layer_collection: A tensorflow_kfac.LayerCollection. Only used by the KFAC
optimizer. Default is None.
recurrent_memory_by_layer: Optional dict, mapping layer names to instances
of transformer_memory.RecurrentMemory. Default is None.
chunk_number: an optional integer Tensor with shape [batch] used to operate
the recurrent_memory.
Returns:
y: a Tensors
"""
x = decoder_input
if hparams.sparse_attention_mode == "sparse":
# If we want to run with our actual sparse kernels, intercept
# the self_attention_type and replace it with our attention fn.
seqlen = common_layers.shape_list(x)[1]
sparse_attention_topology = sparse_matrix.SparseTopology(
"sparse_attention", [seqlen, seqlen],
connector=connectors.Uniform(0.955411645)) # 0.955411659
hparams.self_attention_type = functools.partial(
hparams.self_attention_type, topology=sparse_attention_topology)
elif hparams.sparse_attention_mode == "masked":
# If we're training with sparse attention, create the per-layer
# attention bias that describes the sparsity pattern.
#
# NOTE: We share the same pattern across all attention heads
# within a layer due to memory constraints (because we're not
# actually training with sparse kernels). Per-head patterns
# would likely perform better.
#
# NOTE: We also share the same pattern across all layers, as
# protobuf can't save all of these large tensors if we create
# more than one of them.
decoder_self_attention_bias = generate_sparse_attention_mask(
common_layers.shape_list(x)[1], hparams, 0)
tf.logging.info("Generated sparse attention mask.")
elif hparams.sparse_attention_mode == "dense":
# Replace the dot-product attention with our memory efficient
# version.
hparams.self_attention_type = functools.partial(
hparams.self_attention_type, bias=decoder_self_attention_bias)
pass
else:
# For training on TPU, use T2T's standard attention.
assert hparams.sparse_attention_mode is None
with tf.variable_scope(name):
for layer_idx in range(hparams.num_decoder_layers
or hparams.num_hidden_layers):
x = transformer.transformer_decoder_layer(
x,
decoder_self_attention_bias,
layer_idx,
hparams,
encoder_decoder_attention_bias=encoder_decoder_attention_bias,
encoder_output=encoder_output,
cache=cache,
decode_loop_step=decode_loop_step,
nonpadding=nonpadding,
save_weights_to=save_weights_to,
make_image_summary=make_image_summary,
losses=losses,
layer_collection=layer_collection,
recurrent_memory_by_layer=recurrent_memory_by_layer,
chunk_number=chunk_number)
# if normalization is done in layer_preprocess, then it should also be done
# on the output, since the output can grow very large, being the sum of
# a whole stack of unnormalized layer outputs.
return common_layers.layer_preprocess(
x, hparams, layer_collection=layer_collection)