def testFusedSoftmax(self, m, n):
# Helpers to set up the matrices.
initializer = initializers.Uniform()
# Numpy matrix for verification.
shape = [m, n]
matrix_np = initializer(shape)
# TensorFlow graph.
matrix = tf.Variable(matrix_np, dtype=tf.float32)
output = ops.fused_softmax(matrix)
with self.test_session(use_gpu=True) as sess:
sess.run(tf.global_variables_initializer())
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)
expected_output = softmax(matrix_np)
actual_output = sess.run(output)
self.assertAllClose(actual_output,
expected_output,
atol=1e-03,
rtol=1e-05)
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 testSparseSoftmax(self, m, n, sparsity):
# Helpers to set up the matrices.
connector = connectors.Uniform(sparsity)
initializer = initializers.Uniform()
# Numpy matrix for verification.
matrix_np = connector(initializer([m, n]))
# TensorFlow graph.
matrix = sparse_matrix.SparseMatrix("input", matrix=matrix_np)
output = ops.sparse_softmax(matrix)
with self.test_session(use_gpu=True) as sess:
sess.run(tf.global_variables_initializer())
# 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)
expected_output = self.dense_to_scipy(softmax(matrix_np))
actual_output = self.sparse_to_scipy(
*sess.run(
[output.values, output.row_offsets,
output.column_indices]), expected_output.shape)
self.assert_sparse_matrix_equal(actual_output,
expected_output,
atol=1e-03,
rtol=1e-05)
def testSddmm(self, 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([m, k])
rhs_np = initializer([n, k])
output_np = connector(np.ones([m, n]))
# TensorFlow graph.
output_topology = sparse_matrix.SparseMatrix("output",
matrix=output_np)
lhs = tf.Variable(lhs_np, dtype=tf.float32)
rhs = tf.Variable(rhs_np, dtype=tf.float32)
output = ops.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())
expected_output = self.dense_to_scipy(
output_np * np.dot(lhs_np, np.transpose(rhs_np)))
actual_output = self.sparse_to_scipy(*sess.run(
[output.values, output.row_offsets, output.column_indices]),
shape=expected_output.shape)
self.assert_sparse_matrix_equal(actual_output,
expected_output,
atol=1e-03,
rtol=1e-05)
def __init__(self,
name,
shape=None,
matrix=None,
initializer=initializers.Uniform(),
connector=connectors.Uniform(0.8),
trainable=True,
dtype=tf.float32):
if matrix is None:
assert shape is not None and len(shape) == 2
matrix = connector(initializer(shape))
self._shape = shape
else:
assert shape is None
assert len(matrix.shape) == 2
self._shape = matrix.shape
self._name = name
self._trainable = trainable
self._dtype = dtype
self._sparsity = 1.0 - np.count_nonzero(matrix) / matrix.size
# Create a numpy version of the sparse matrix.
values_, row_indices_, row_offsets_, column_indices_ = _dense_to_sparse(
matrix)
# Create tensors for the matrix shape on the host. These are for internal
# use and should generally not be used by end-user. Use the normal python
# 'shape' property instead.
with tf.device("cpu"):
self._rows = tf.get_variable(initializer=self._shape[0],
trainable=False,
name=self._name + "_rows",
dtype=tf.int32)
self._columns = tf.get_variable(initializer=self._shape[1],
trainable=False,
name=self._name + "_columns",
dtype=tf.int32)
# Convert the sparse matrix to TensorFlow variables.
self._values = tf.get_variable(initializer=values_,
trainable=self.trainable,
name=self._name + "_values",
dtype=self._dtype)
self._row_indices = tf.get_variable(initializer=row_indices_,
trainable=False,
name=self._name + "_row_indices",
dtype=tf.uint32)
self._row_offsets = tf.get_variable(initializer=row_offsets_,
trainable=False,
name=self._name + "_row_offsets",
dtype=tf.uint32)
self._column_indices = tf.get_variable(initializer=column_indices_,
trainable=False,
name=self._name +
"_column_indices",
dtype=tf.uint32)
# Add this matrix to the collection of trainable matrices.
track_trainable_sparse_matrix(self)
def testSpmmGradient(self, 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 = connector(initializer([m, k]))
rhs_np = initializer([k, n])
lhs = sparse_matrix.SparseMatrix("lhs", matrix=lhs_np)
rhs = tf.Variable(rhs_np, dtype=tf.float32)
output = ops.spmm(lhs, rhs)
with self.test_session(use_gpu=use_gpu) as sess:
sess.run(tf.global_variables_initializer())
error = tf.test.compute_gradient_error(
[lhs.values, rhs], [lhs.values.shape.as_list(), [k, n]],
output, [m, n])
self.assertLess(error, 1e-3)
def testSpmm(self, 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 = connector(initializer([m, k]))
rhs_np = initializer([k, n])
# TensorFlow graph.
lhs = sparse_matrix.SparseMatrix("lhs", matrix=lhs_np)
rhs = tf.Variable(rhs_np, dtype=tf.float32)
output = ops.spmm(lhs, rhs)
# Execute the op and compare the results.
with self.test_session(use_gpu=use_gpu) as sess:
sess.run(tf.global_variables_initializer())
self.assertAllClose(sess.run(output),
np.dot(lhs_np, rhs_np),
atol=1e-03,
rtol=1e-05)
def testSddmm_Replicated(self, r, m, k, n, sparsity, force_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(force_gpu=force_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 testCsr2Idx(self, m, n, sparsity, use_gpu):
# Helpers to set up the matrices.
connector = connectors.Uniform(sparsity)
initializer = initializers.Uniform()
# Numpy matrix for verification.
matrix_np = connector(initializer([m, n]))
# TensorFlow graph.
matrix = sparse_matrix.SparseMatrix("input", matrix=matrix_np)
output = ops.csr2idx(matrix)
# Execute the op and compare the results.
with self.test_session(use_gpu=use_gpu) as sess:
sess.run(tf.global_variables_initializer())
# Calculate the linear indices in numpy.
x = self.dense_to_scipy(matrix_np)
expected_output = np.concatenate([
x.indices[x.indptr[i]:x.indptr[i + 1]] + i * n
for i in range(m)
])
self.assertAllEqual(sess.run(output), expected_output)
def testTranspose(self, m, n, sparsity, use_gpu):
# Helpers to set up the matrices.
connector = connectors.Uniform(sparsity)
initializer = initializers.Uniform()
# Numpy matrix for verification.
matrix_np = connector(initializer([m, n]))
# TensorFlow graph.
matrix = sparse_matrix.SparseMatrix("input", matrix=matrix_np)
output = ops.transpose(matrix)
# Execute the op and compare the results.
with self.test_session(use_gpu=use_gpu) as sess:
sess.run(tf.global_variables_initializer())
expected_output = self.dense_to_scipy(np.transpose(matrix_np))
actual_output = self.sparse_to_scipy(*sess.run(
[output.values, output.row_offsets, output.column_indices]),
shape=expected_output.shape)
self.assert_sparse_matrix_equal(actual_output,
expected_output,
atol=1e-03,
rtol=1e-05)
def testSddmmGradient(self, 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([m, k])
rhs_np = initializer([n, k])
output_np = connector(np.ones([m, n]))
# TensorFlow graph.
output_topology = sparse_matrix.SparseMatrix("output",
matrix=output_np)
lhs = tf.Variable(lhs_np, dtype=tf.float32)
rhs = tf.Variable(rhs_np, dtype=tf.float32)
output = ops.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())
error = tf.test.compute_gradient_error(
[lhs, rhs], [[m, k], [n, k]], output.values,
output.values.shape.as_list())
self.assertLess(error, 1e-3)
def testDenseToSparse(self, m, n, sparsity):
# Helpers to set up the matrices.
connector = connectors.Uniform(sparsity)
initializer = initializers.Uniform()
# Create a dense matrix in numpy with the specified sparsity.
matrix = connector(initializer([m, n]))
# Convert to a sparse numpy matrix.
values, row_indices, row_offsets, column_indices = sparse_matrix._dense_to_sparse(
matrix)
# Create a scipy version of the matrix.
expected_output = scipy.sparse.csr_matrix(
(values, column_indices, row_offsets), [m, n])
# Create the expected row indices.
expected_row_indices = np.argsort(-1 * np.diff(expected_output.indptr))
# Compare the matrices.
self.assertAllEqual(expected_output.data, values)
self.assertAllEqual(expected_output.indptr, row_offsets)
self.assertAllEqual(expected_output.indices, column_indices)
self.assertAllEqual(expected_row_indices, row_indices)
