def benchmarkTridiagonalMulOp(self):
devices = [('/cpu:0', 'cpu')]
for device_id, device_name in devices:
for batch_size, matrix_size in self.sizes:
with ops.Graph().as_default(), \
session.Session(config=benchmark.benchmark_config()) as sess, \
ops.device(device_id):
upper, diag, lower, vec = self._generateData(
batch_size, matrix_size)
x1 = self.baseline(upper, diag, lower, vec)
x2 = linalg_impl.tridiagonal_matmul(
(upper, diag, lower), vec)
variables.global_variables_initializer().run()
self.run_op_benchmark(
sess,
control_flow_ops.group(x1),
min_iters=10,
store_memory_usage=False,
name=('tridiagonal_matmul_baseline_%s'
'_batch_size_%d_matrix_size_%d' %
(device_name, batch_size, matrix_size)))
self.run_op_benchmark(
sess,
control_flow_ops.group(x2),
min_iters=10,
store_memory_usage=False,
name=
('tridiagonal_matmul_%s_batch_size_%d_matrix_size_%d'
% (device_name, batch_size, matrix_size)))
def benchmarkTridiagonalMulOp(self):
devices = [('/cpu:0', 'cpu')]
if test.is_gpu_available(cuda_only=True):
devices += [('/gpu:0', 'gpu')]
for device_option, size_option in itertools.product(devices, self.sizes):
device_id, device_name = device_option
m, batch_size, n = size_option
with ops.Graph().as_default(), \
session.Session(config=benchmark.benchmark_config()) as sess, \
ops.device(device_id):
upper, diag, lower, vec = self._generateData(batch_size, m, n)
x1 = self.baseline(upper, diag, lower, vec)
x2 = linalg_impl.tridiagonal_matmul((upper, diag, lower),
vec,
diagonals_format='sequence')
variables.global_variables_initializer().run()
self.run_op_benchmark(
sess,
control_flow_ops.group(x1),
min_iters=10,
store_memory_usage=False,
name=('tridiagonal_matmul_baseline_%s'
'_batch_size_%d_m_%d_n_%d' %
(device_name, batch_size, m, n)))
self.run_op_benchmark(
sess,
control_flow_ops.group(x2),
min_iters=10,
store_memory_usage=False,
name=('tridiagonal_matmul_%s_batch_size_%d_m_%d_n_%d' %
(device_name, batch_size, m, n)))
def _matmul(self, x, adjoint=False, adjoint_arg=False):
diagonals = self.diagonals
if adjoint:
diagonals = self._construct_adjoint_diagonals(diagonals)
x = linalg.adjoint(x) if adjoint_arg else x
return linalg.tridiagonal_matmul(
diagonals, x, diagonals_format=self.diagonals_format)
def _testAllFormats(self,
superdiag,
maindiag,
subdiag,
rhs,
expected,
dtype=dtypes.float64):
superdiag_extended = np.pad(superdiag, [0, 1], 'constant')
subdiag_extended = np.pad(subdiag, [1, 0], 'constant')
diags_compact = np.stack([superdiag_extended, maindiag, subdiag_extended])
diags_matrix = np.diag(superdiag, 1) + np.diag(maindiag, 0) + np.diag(
subdiag, -1)
diags_sequence = (constant_op.constant(superdiag_extended, dtype),
constant_op.constant(maindiag, dtype),
constant_op.constant(subdiag_extended, dtype))
diags_compact = constant_op.constant(diags_compact, dtype)
diags_matrix = constant_op.constant(diags_matrix, dtype)
rhs = constant_op.constant(rhs, dtype)
rhs_batch = array_ops.stack([rhs, 2 * rhs])
diags_compact_batch = array_ops.stack([diags_compact, 2 * diags_compact])
diags_matrix_batch = array_ops.stack([diags_matrix, 2 * diags_matrix])
diags_sequence_batch = [array_ops.stack([x, 2 * x]) for x in diags_sequence]
results = [
linalg_impl.tridiagonal_matmul(
diags_sequence, rhs, diagonals_format='sequence'),
linalg_impl.tridiagonal_matmul(
diags_compact, rhs, diagonals_format='compact'),
linalg_impl.tridiagonal_matmul(
diags_matrix, rhs, diagonals_format='matrix')
]
results_batch = [
linalg_impl.tridiagonal_matmul(
diags_sequence_batch, rhs_batch, diagonals_format='sequence'),
linalg_impl.tridiagonal_matmul(
diags_compact_batch, rhs_batch, diagonals_format='compact'),
linalg_impl.tridiagonal_matmul(
diags_matrix_batch, rhs_batch, diagonals_format='matrix')
]
with self.cached_session(use_gpu=True):
results = self.evaluate(results)
results_batch = self.evaluate(results_batch)
expected = np.array(expected)
expected_batch = np.stack([expected, 4 * expected])
for result in results:
self.assertAllClose(result, expected)
for result in results_batch:
self.assertAllClose(result, expected_batch)
def testBatch(self):
b = 20
m = 10
n = 15
superdiag = self._randomComplexArray((b, m - 1))
maindiag = self._randomComplexArray((b, m))
subdiag = self._randomComplexArray((b, m - 1))
rhs = self._randomComplexArray((b, m, n))
matrix = np.stack([np.diag(superdiag[i], 1) + \
np.diag(maindiag[i], 0) + \
np.diag(subdiag[i], -1) for i in range(b)])
expected_result = np.matmul(matrix, rhs)
result = linalg_impl.tridiagonal_matmul(
constant_op.constant(matrix, dtype=dtypes.complex128),
constant_op.constant(rhs, dtype=dtypes.complex128),
diagonals_format='matrix')
with self.cached_session():
result = self.evaluate(result)
self.assertAllClose(result, expected_result)
def testBatch(self):
b = 20
m = 10
n = 15
superdiag = self._randomComplexArray((b, m - 1))
maindiag = self._randomComplexArray((b, m))
subdiag = self._randomComplexArray((b, m - 1))
rhs = self._randomComplexArray((b, m, n))
matrix = np.stack([np.diag(superdiag[i], 1) + \
np.diag(maindiag[i], 0) + \
np.diag(subdiag[i], -1) for i in range(b)])
expected_result = np.matmul(matrix, rhs)
result = linalg_impl.tridiagonal_matmul(
constant_op.constant(matrix, dtype=dtypes.complex128),
constant_op.constant(rhs, dtype=dtypes.complex128),
diagonals_format='matrix')
with self.cached_session(use_gpu=True):
result = self.evaluate(result)
self.assertAllClose(result, expected_result)
def benchmarkTridiagonalMulOp(self):
devices = [('/cpu:0', 'cpu')]
if test.is_gpu_available(cuda_only=True):
devices += [('/gpu:0', 'gpu')]
for device_option, size_option in itertools.product(
devices, self.sizes):
device_id, device_name = device_option
m, batch_size, n = size_option
with ops.Graph().as_default(), \
session.Session(config=benchmark.benchmark_config()) as sess, \
ops.device(device_id):
upper, diag, lower, vec = self._generateData(
batch_size, m, n)
x1 = self.baseline(upper, diag, lower, vec)
x2 = linalg_impl.tridiagonal_matmul(
(upper, diag, lower), vec, diagonals_format='sequence')
self.evaluate(variables.global_variables_initializer())
self.run_op_benchmark(
sess,
control_flow_ops.group(x1),
min_iters=10,
store_memory_usage=False,
name=('tridiagonal_matmul_baseline_%s'
'_batch_size_%d_m_%d_n_%d' %
(device_name, batch_size, m, n)))
self.run_op_benchmark(
sess,
control_flow_ops.group(x2),
min_iters=10,
store_memory_usage=False,
name=('tridiagonal_matmul_%s_batch_size_%d_m_%d_n_%d' %
(device_name, batch_size, m, n)))
def _testAllFormats(self,
superdiag,
maindiag,
subdiag,
rhs,
expected,
dtype=dtypes.float64):
superdiag_extended = np.pad(superdiag, [0, 1], 'constant')
subdiag_extended = np.pad(subdiag, [1, 0], 'constant')
diags_compact = np.stack(
[superdiag_extended, maindiag, subdiag_extended])
diags_matrix = np.diag(superdiag, 1) + np.diag(maindiag, 0) + np.diag(
subdiag, -1)
diags_sequence = (constant_op.constant(superdiag_extended, dtype),
constant_op.constant(maindiag, dtype),
constant_op.constant(subdiag_extended, dtype))
diags_compact = constant_op.constant(diags_compact, dtype)
diags_matrix = constant_op.constant(diags_matrix, dtype)
rhs = constant_op.constant(rhs, dtype)
rhs_batch = array_ops.stack([rhs, 2 * rhs])
diags_compact_batch = array_ops.stack(
[diags_compact, 2 * diags_compact])
diags_matrix_batch = array_ops.stack([diags_matrix, 2 * diags_matrix])
diags_sequence_batch = [
array_ops.stack([x, 2 * x]) for x in diags_sequence
]
results = [
linalg_impl.tridiagonal_matmul(diags_sequence,
rhs,
diagonals_format='sequence'),
linalg_impl.tridiagonal_matmul(diags_compact,
rhs,
diagonals_format='compact'),
linalg_impl.tridiagonal_matmul(diags_matrix,
rhs,
diagonals_format='matrix')
]
results_batch = [
linalg_impl.tridiagonal_matmul(diags_sequence_batch,
rhs_batch,
diagonals_format='sequence'),
linalg_impl.tridiagonal_matmul(diags_compact_batch,
rhs_batch,
diagonals_format='compact'),
linalg_impl.tridiagonal_matmul(diags_matrix_batch,
rhs_batch,
diagonals_format='matrix')
]
with self.cached_session(use_gpu=False):
results = self.evaluate(results)
results_batch = self.evaluate(results_batch)
expected = np.array(expected)
expected_batch = np.stack([expected, 4 * expected])
for result in results:
self.assertAllClose(result, expected)
for result in results_batch:
self.assertAllClose(result, expected_batch)
def _tridiagonal_matmul(self, diagonals, rhs, diagonals_format):
return linalg_impl.tridiagonal_matmul(
diagonals, rhs, diagonals_format=diagonals_format)
