def sparse_matmul(self, x, feature_axis, output_dim):
"""
:param tf.Tensor x:
:param int feature_axis:
:param int output_dim:
:return: y, weights, bsmm
:rtype: (tf.Tensor, tf.Variable, object)
"""
block_size = self.block_size
input_dim = x.get_shape().dims[feature_axis].value
assert input_dim is not None, "%r shape unknown" % (x,)
assert input_dim % block_size == 0 and output_dim % block_size == 0
from blocksparse.matmul import BlocksparseMatMul
seed = self.random.randint(2 ** 31)
sparsity_pattern = sparsity_pattern_barabasi_albert(
n1=input_dim // block_size, n2=output_dim // block_size, m=self.connectivity, dense=self.connectivity_dense,
seed=seed)
bsmm = BlocksparseMatMul(sparsity_pattern, block_size=block_size, feature_axis=feature_axis)
if self.weights_identity_init:
weights_init = bsmm.identity_init()
else:
weights_init = None
weights = tf.get_variable("W", shape=bsmm.w_shape, initializer=weights_init)
y = bsmm(x, weights)
return y, weights, bsmm
def test_blocksparse_simple_feature_axis1():
init_blocksparse()
from blocksparse.matmul import BlocksparseMatMul
import tensorflow as tf
import numpy
n_in = 64
n_out = 32 * 32
block_size = 32
n_batch = 4
# Create a dense sparsity pattern
mask = numpy.ones((n_in // block_size, n_out // block_size), dtype=numpy.int32)
# MatMul object
bsmm = BlocksparseMatMul(mask, block_size=block_size, feature_axis=1, name="bsmm")
# Input
x_np = numpy.arange(n_in * n_batch, dtype=numpy.float32).reshape((n_batch, n_in)) + 1.0
x = tf.constant(x_np, name='x')
# Block-sparse weights
w_np = bsmm.identity_init()()
w = tf.constant(w_np, name="w")
# Block-sparse matrix multiplication
y = bsmm(x, w)
y.set_shape((n_batch, n_out))
# Run
result = session.run(y)
print(result)
print('L2:', numpy.sum(result ** 2))
y_test = bsmm.fprop_test(x_np, w_np)
print(y_test)
assert_allclose(result, y_test)
def test_blocksparse_simple_identity():
init_blocksparse()
from blocksparse.matmul import BlocksparseMatMul
import tensorflow as tf
import numpy
n_in = 64
n_out = 32 * 32
block_size = 32
# Note: It seems everything less than 4 fails, as well as non-power-of-2.
n_batch = 4
# Create a dense sparsity pattern
mask = numpy.ones((n_in // block_size, n_out // block_size),
dtype=numpy.int32)
# MatMul object
bsmm = BlocksparseMatMul(mask,
block_size=block_size,
feature_axis=0,
name="bsmm")
# Input
x_np = numpy.arange(n_in * n_batch, dtype=numpy.float32).reshape(
(n_in, n_batch)) + 1.0
x = tf.constant(x_np, name='x')
# Block-sparse weights
w_np = bsmm.identity_init()()
w = tf.constant(w_np, name="w")
#for b in range(bsmm.blocks):
# cb, kb = bsmm.updat_list[b]
# print("block %i/%i, cb %i/%i, kb %i/%i" % (b, bsmm.blocks, cb, bsmm.KB, kb, bsmm.CB))
# Block-sparse matrix multiplication
y = bsmm(x, w)
y.set_shape((n_out, n_batch))
# Run
result = session.run(y)
print(result)
print('L2:', numpy.sum(result**2))
y_test = bsmm.fprop_test(x_np, w_np)
print(y_test)
i = numpy.argmax((y_test - result)**2)
print('biggest diff at %i: %r vs %r' %
(i, y_test.flatten()[i], result.flatten()[i]))
assert_allclose(result, y_test, rtol=1e-2)