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