1# Copyright 2015 The TensorFlow Authors. All Rights Reserved. 2# 3# Licensed under the Apache License, Version 2.0 (the "License"); 4# you may not use this file except in compliance with the License. 5# You may obtain a copy of the License at 6# 7# http://www.apache.org/licenses/LICENSE-2.0 8# 9# Unless required by applicable law or agreed to in writing, software 10# distributed under the License is distributed on an "AS IS" BASIS, 11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12# See the License for the specific language governing permissions and 13# limitations under the License. 14# ============================================================================== 15"""Tests for sparse_ops.sparse_tensor_dense_matmul.""" 16 17import sys 18import time 19 20from absl import app 21import numpy as np 22 23from tensorflow.core.protobuf import config_pb2 24from tensorflow.python.client import session 25from tensorflow.python.framework import constant_op 26from tensorflow.python.framework import dtypes 27from tensorflow.python.framework import ops 28from tensorflow.python.framework import sparse_tensor 29from tensorflow.python.framework import tensor_shape 30from tensorflow.python.framework import test_util 31from tensorflow.python.ops import array_ops 32from tensorflow.python.ops import control_flow_ops 33from tensorflow.python.ops import math_ops 34from tensorflow.python.ops import sparse_ops 35from tensorflow.python.platform import test 36 37 38def _maybe_complex(x): 39 if x.dtype.kind == "c": # complex 40 return x + 1j * x 41 return x 42 43 44class SparseTensorDenseMatMulTest(test.TestCase): 45 46 def _testMatmul(self, 47 x, 48 y, 49 adjoint_a=False, 50 adjoint_b=False, 51 indices_dtype=np.int64): 52 x_mat = np.array(x) 53 if adjoint_a: 54 x_mat = x_mat.T.conj() 55 y_mat = np.array(y) 56 if adjoint_b: 57 y_mat = y_mat.T.conj() 58 59 np_ans = x_mat.dot(y_mat) 60 61 x_indices = np.vstack(np.where(x)).astype(indices_dtype).T 62 x_values = x[np.where(x)] 63 x_shape = x.shape 64 65 with self.cached_session(): 66 sp_x_value = sparse_tensor.SparseTensorValue( 67 indices=x_indices, values=x_values, dense_shape=x_shape) 68 tf_value_ans = sparse_ops.sparse_tensor_dense_matmul( 69 sp_x_value, y, adjoint_a=adjoint_a, adjoint_b=adjoint_b) 70 tf_tensor_ans = sparse_ops.sparse_tensor_dense_matmul( 71 sparse_tensor.SparseTensor.from_value(sp_x_value), 72 y, 73 adjoint_a=adjoint_a, 74 adjoint_b=adjoint_b) 75 76 # Ensure that the RHS shape is known at least. 77 self.assertEqual(tf_value_ans.get_shape()[1], np_ans.shape[1]) 78 self.assertEqual(tf_tensor_ans.get_shape()[1], np_ans.shape[1]) 79 80 for out in (self.evaluate(tf_value_ans), self.evaluate(tf_tensor_ans)): 81 if x.dtype == np.float32: 82 self.assertAllClose(np_ans, out, rtol=1e-4, atol=1e-4) 83 elif x.dtype == np.float64: 84 self.assertAllClose(np_ans, out, rtol=1e-6, atol=1e-6) 85 elif x.dtype == np.float16: 86 self.assertAllClose(np_ans, out, rtol=1e-3, atol=1e-3) 87 else: 88 self.assertAllClose(np_ans, out, rtol=1e-3, atol=1e-3) 89 90 def _testBasic(self, value_dtype, indices_dtype=np.int64): 91 x = _maybe_complex(np.random.rand(10, 10).astype(value_dtype)) 92 x[np.abs(x) < 0.5] = 0 # Make it sparse 93 94 y = _maybe_complex(np.random.randn(10, 20).astype(value_dtype)) 95 96 self._testMatmul(x, y, indices_dtype=indices_dtype) 97 98 def testBasic(self): 99 np.random.seed(127) # Repeatable results 100 self._testBasic(np.int32) 101 self._testBasic(np.float16) 102 self._testBasic(np.float32) 103 self._testBasic(np.float64) 104 self._testBasic(np.complex64) 105 self._testBasic(np.complex128) 106 self._testBasic(np.int32, indices_dtype=np.int32) 107 self._testBasic(np.float32, indices_dtype=np.int32) 108 109 def testShapeInference(self): 110 x = np.random.rand(10, 10) 111 x[np.abs(x) < 0.5] = 0 # Make it sparse 112 y = np.random.randn(10, 20) 113 x_indices = np.vstack(np.where(x)).astype(np.int64).T 114 x_values = x[np.where(x)] 115 x_shape = x.shape 116 117 with ops.Graph().as_default(): 118 x_st = sparse_tensor.SparseTensor(x_indices, x_values, x_shape) 119 result = sparse_ops.sparse_tensor_dense_matmul(x_st, y) 120 self.assertEqual(result.get_shape(), (10, 20)) 121 122 x_shape_unknown = array_ops.placeholder(dtype=dtypes.int64, shape=None) 123 x_st_shape_unknown = sparse_tensor.SparseTensor(x_indices, x_values, 124 x_shape_unknown) 125 result_left_shape_unknown = sparse_ops.sparse_tensor_dense_matmul( 126 x_st_shape_unknown, y) 127 self.assertEqual(result_left_shape_unknown.get_shape().as_list(), 128 [None, 20]) 129 130 x_shape_inconsistent = [10, 15] 131 x_st_shape_inconsistent = sparse_tensor.SparseTensor( 132 x_indices, x_values, x_shape_inconsistent) 133 with self.assertRaisesRegex(ValueError, "Dimensions must be equal"): 134 sparse_ops.sparse_tensor_dense_matmul(x_st_shape_inconsistent, y) 135 136 @test_util.run_in_graph_and_eager_modes(use_gpu=False) 137 def testInvalidIndicesForSparseTensorDenseMatmul(self): 138 # TODO(b/169813429): Make GPU kernel return nice errors too. 139 indices = np.array([[1, 10]]).astype(np.int64) 140 values = np.array([10]).astype(np.float32) 141 shape = [3, 2] 142 sparse_t = sparse_tensor.SparseTensor(indices, values, shape) 143 144 # Test multiplying by both a small and large dense matrix, to hit 145 # both cases in the kernel. 146 dense_t = np.array([[1] * 5, [2] * 5], dtype=np.float32) 147 with self.assertRaisesOpError("k .10. from index.0,1. out of bounds .>=2."): 148 self.evaluate(sparse_ops.sparse_tensor_dense_matmul(sparse_t, dense_t)) 149 dense_t = np.array([[1] * 500, [2] * 500], dtype=np.float32) 150 with self.assertRaisesOpError("k .10. from index.0,1. out of bounds .>=2."): 151 self.evaluate(sparse_ops.sparse_tensor_dense_matmul(sparse_t, dense_t)) 152 153 # Repeat with adjoint_a, to get a different error. 154 dense_t = np.array([[1] * 5, [2] * 5, [3] * 5], dtype=np.float32) 155 with self.assertRaisesOpError("m .10. from index.0,1. out of bounds .>=2."): 156 self.evaluate( 157 sparse_ops.sparse_tensor_dense_matmul( 158 sparse_t, dense_t, adjoint_a=True)) 159 dense_t = np.array([[1] * 500, [2] * 500, [3] * 500], dtype=np.float32) 160 with self.assertRaisesOpError("m .10. from index.0,1. out of bounds .>=2."): 161 self.evaluate( 162 sparse_ops.sparse_tensor_dense_matmul( 163 sparse_t, dense_t, adjoint_a=True)) 164 165 def testUnorderedIndicesForSparseTensorDenseMatmul(self): 166 indices = np.array([(2, 1), (0, 0)]).astype(np.int64) 167 values = np.array([10, 11]).astype(np.float32) 168 shape = [3, 2] 169 sparse_t = sparse_tensor.SparseTensor(indices, values, shape) 170 171 dense_t = np.array([[1] * 500, [2] * 500], dtype=np.float32) 172 expected_t = np.array([[11] * 500, [0] * 500, [20] * 500], dtype=np.float32) 173 174 self.assertAllClose( 175 expected_t, sparse_ops.sparse_tensor_dense_matmul(sparse_t, dense_t)) 176 177 @test_util.run_gpu_only 178 def testInvalidIndicesForSparseTensorDenseMatmulOnGPU(self): 179 indices = np.array([[1, 10]]).astype(np.int64) 180 values = np.array([10]).astype(np.float32) 181 shape = [3, 2] 182 sparse_t = sparse_tensor.SparseTensor(indices, values, shape) 183 184 # Test multiplying by both a small and large dense matrix, to hit 185 # both cases in the kernel. 186 dense_t = np.array([[1] * 5, [2] * 5], dtype=np.float32) 187 expected_t = np.array([[0] * 5, [np.nan] * 5, [0] * 5], dtype=np.float32) 188 self.assertAllClose( 189 expected_t, sparse_ops.sparse_tensor_dense_matmul(sparse_t, dense_t)) 190 dense_t = np.array([[1] * 500, [2] * 500], dtype=np.float32) 191 expected_t = np.array([[0] * 500, [np.nan] * 500, [0] * 500], 192 dtype=np.float32) 193 self.assertAllClose( 194 expected_t, sparse_ops.sparse_tensor_dense_matmul(sparse_t, dense_t)) 195 196 # Repeat with adjoint_a, now the error is that the sparse index 197 # is OOO w.r.t. the output. The GPU kernel can't do much here, 198 # so it just doesn't accumulate. 199 200 dense_t = np.array([[1] * 5, [2] * 5, [3] * 5], dtype=np.float32) 201 expected_t = np.array([[0] * 5, [0] * 5], dtype=np.float32) 202 self.assertAllClose( 203 expected_t, 204 sparse_ops.sparse_tensor_dense_matmul( 205 sparse_t, dense_t, adjoint_a=True)) 206 207 dense_t = np.array([[1] * 500, [2] * 500, [3] * 500], dtype=np.float32) 208 expected_t = np.array([[0] * 500, [0] * 500], dtype=np.float32) 209 self.assertAllClose( 210 expected_t, 211 sparse_ops.sparse_tensor_dense_matmul( 212 sparse_t, dense_t, adjoint_a=True)) 213 214 def _testLarge(self, np_dtype): 215 r1 = np.random.randint(6000, 20000) 216 r2 = np.random.randint(1, 10) 217 r3 = np.random.randint(1, 10) 218 219 for m, k, n in [(r1, r2, r3), 220 (r2, r1, r3), 221 (r2, r3, r1)]: 222 x = _maybe_complex(np.random.rand(m, k).astype(np_dtype)) 223 x[np.abs(x) < 0.8] = 0 224 225 y = _maybe_complex(np.random.randn(k, n).astype(np_dtype)) 226 227 self._testMatmul(x, y, adjoint_a=False, adjoint_b=False) 228 self._testMatmul(x.transpose(), y, adjoint_a=True, adjoint_b=False) 229 self._testMatmul(x, y.transpose(), adjoint_a=False, adjoint_b=True) 230 self._testMatmul( 231 x.transpose(), y.transpose(), adjoint_a=True, adjoint_b=True) 232 233 # Tests setting one dimension to be a high value. 234 def testLarge(self): 235 np.random.seed(127) # Repeatable results 236 self._testLarge(np.float32) 237 self._testLarge(np.float64) 238 self._testLarge(np.complex64) 239 self._testLarge(np.complex128) 240 241 # Tests random sized matrices. 242 def testFloatRandom(self): 243 np.random.seed(127) # Repeatable results 244 for _ in range(8): 245 for adjoint_a in [True, False]: 246 for adjoint_b in [True, False]: 247 for thresh in [0.0, 0.2, 0.8, 1.0]: 248 n, k, m = np.random.randint(1, 100, size=3) 249 x = np.random.rand(n, k).astype(np.float32) 250 x[x < thresh] = 0 # Make it sparse 251 y = np.random.randn(k, m).astype(np.float32) 252 x = x.transpose() if adjoint_a else x 253 y = y.transpose() if adjoint_b else y 254 self._testMatmul(x, y, adjoint_a, adjoint_b) 255 256 257def _sparse_tensor_dense_vs_dense_matmul_benchmark_dense(x, y, adjoint_a, 258 adjoint_b): 259 260 def body(t, prev): 261 with ops.control_dependencies([prev]): 262 return (t + 1, math_ops.matmul( 263 x, 264 y, 265 transpose_a=adjoint_a, 266 transpose_b=adjoint_b, 267 a_is_sparse=True, 268 b_is_sparse=False)) 269 270 t0 = constant_op.constant(0) 271 v0 = constant_op.constant(0.0) 272 273 def _timeit(iterations, _): 274 (_, final) = control_flow_ops.while_loop( 275 lambda t, _: t < iterations, 276 body, (t0, v0), 277 parallel_iterations=1, 278 back_prop=False, 279 shape_invariants=(tensor_shape.TensorShape(()), 280 tensor_shape.TensorShape(None))) 281 return [final] 282 283 return _timeit 284 285 286def _sparse_tensor_dense_vs_dense_matmul_benchmark_sparse(x_ind, x_val, x_shape, 287 y, adjoint_a, 288 adjoint_b): 289 sp_x = sparse_tensor.SparseTensor( 290 indices=x_ind, values=x_val, dense_shape=x_shape) 291 292 def body(t, prev): 293 with ops.control_dependencies([prev]): 294 return (t + 1, sparse_ops.sparse_tensor_dense_matmul( 295 sp_x, y, adjoint_a=adjoint_a, adjoint_b=adjoint_b)) 296 297 t0 = constant_op.constant(0) 298 v0 = constant_op.constant(0.0) 299 300 def _timeit(iterations, _): 301 (_, final) = control_flow_ops.while_loop( 302 lambda t, _: t < iterations, 303 body, (t0, v0), 304 parallel_iterations=1, 305 back_prop=False, 306 shape_invariants=(tensor_shape.TensorShape(()), 307 tensor_shape.TensorShape(None))) 308 return [final] 309 310 return _timeit 311 312 313def sparse_tensor_dense_vs_dense_matmul_benchmark(thresh, 314 m, 315 k, 316 n, 317 adjoint_a, 318 adjoint_b, 319 use_gpu, 320 skip_dense=False): 321 config = config_pb2.ConfigProto() 322 config.allow_soft_placement = True 323 324 # Configurable for benchmarking: 325 # config.intra_op_parallelism_threads = 100 326 # config.gpu_options.per_process_gpu_memory_fraction = 0.3 327 328 np.random.seed([6, 117]) # Reproducibility 329 x = np.random.rand(m, k).astype(np.float32) 330 x[x < thresh] = 0 331 y = np.random.randn(k, n).astype(np.float32) 332 if adjoint_a: 333 x = x.T 334 if adjoint_b: 335 y = y.T 336 337 def _timer(sess, ops_fn, iterations): 338 # Warm in 339 sess.run(ops_fn(10, sess)) 340 341 # Timing run 342 start = time.time() 343 sess.run(ops_fn(iterations, sess)) 344 end = time.time() 345 346 return (end - start) / (1.0 * iterations) # Average runtime per iteration 347 348 # Using regular matmul, marking one of the matrices as dense. 349 if skip_dense: 350 delta_dense = float("nan") 351 else: 352 with session.Session(config=config, graph=ops.Graph()) as sess: 353 if not use_gpu: 354 with ops.device("/cpu:0"): 355 x_t = constant_op.constant(x) 356 y_t = constant_op.constant(y) 357 ops_fn = _sparse_tensor_dense_vs_dense_matmul_benchmark_dense( 358 x_t, y_t, adjoint_a, adjoint_b) 359 else: 360 with ops.device("/device:GPU:0"): 361 x_t = constant_op.constant(x) 362 y_t = constant_op.constant(y) 363 ops_fn = _sparse_tensor_dense_vs_dense_matmul_benchmark_dense( 364 x_t, y_t, adjoint_a, adjoint_b) 365 delta_dense = _timer(sess, ops_fn, 200) 366 367 # Using sparse_tensor_dense_matmul. 368 with session.Session("", config=config, graph=ops.Graph()) as sess: 369 if not use_gpu: 370 with ops.device("/cpu:0"): 371 x_ind = constant_op.constant(np.vstack(np.where(x)).astype(np.int64).T) 372 x_val = constant_op.constant(x[np.where(x)]) 373 x_shape = constant_op.constant(np.array(x.shape).astype(np.int64)) 374 y_t = constant_op.constant(y) 375 ops_fn = _sparse_tensor_dense_vs_dense_matmul_benchmark_sparse( 376 x_ind, x_val, x_shape, y_t, adjoint_a, adjoint_b) 377 else: 378 with ops.device("/device:GPU:0"): 379 x_ind = constant_op.constant(np.vstack(np.where(x)).astype(np.int64).T) 380 x_val = constant_op.constant(x[np.where(x)]) 381 x_shape = constant_op.constant(np.array(x.shape).astype(np.int64)) 382 y_t = constant_op.constant(y) 383 ops_fn = _sparse_tensor_dense_vs_dense_matmul_benchmark_sparse( 384 x_ind, x_val, x_shape, y_t, adjoint_a, adjoint_b) 385 delta_sparse = _timer(sess, ops_fn, 200) 386 387 print("%g \t %d \t %s \t %d \t %d \t %g \t %g \t %g" % 388 (1 - thresh, n, use_gpu, m, k, delta_dense, delta_sparse, 389 delta_sparse / delta_dense)) 390 391 392def main(_): 393 print("DenseDense MatMul (w/ Sparse Flag) vs. SparseTensorDense MatMul") 394 print("Matrix sizes:") 395 print(" A sparse [m, k] with % nonzero values between 1% and 80%") 396 print(" B dense [k, n]") 397 print("") 398 print("% nnz \t n \t gpu \t m \t k \t dt(dense) \t dt(sparse) " 399 "\t dt(sparse)/dt(dense)") 400 401 for thresh in (0.99, 0.8, 0.5, 0.2): 402 for n in (50, 100): 403 for use_gpu in (True, False): 404 for m in (100, 1000): 405 for k in (100, 1000): 406 sparse_tensor_dense_vs_dense_matmul_benchmark( 407 thresh, m, k, n, False, False, use_gpu=use_gpu) 408 409 # Enable for large scale benchmarks, these ones take a long time to run. 410 # 411 # for use_gpu in (True, False): 412 # sparse_tensor_dense_vs_dense_matmul_benchmark( 413 # thresh=0.99, m=1000000, k=1000, n=100, adjoint_a=False, 414 # adjoint_b=False, use_gpu=use_gpu, skip_dense=True) 415 416 417if __name__ == "__main__": 418 if "--benchmarks" in sys.argv: 419 sys.argv.remove("--benchmarks") 420 app.run() # pylint: disable=no-value-for-parameter 421 else: 422 test.main() 423