def create_sparse_layers(opts):
matmul_opts = {"metaInfoBucketOversizeProportion": opts.meta_info_oversize}
in_blocks = opts.input_size // opts.block_size
out_blocks = opts.output_size // opts.block_size
identity_size = max(in_blocks, out_blocks)
block_mask = np.identity(identity_size)[0:in_blocks, 0:out_blocks]
block_mask[1, 3] = 1
block_mask[0, 3] = 1
n_blocks = np.count_nonzero(block_mask)
el_mask = sparse.block_mask_to_element(block_mask, opts.block_size)
n_els = np.count_nonzero(el_mask)
masked_rhs = np.zeros_like(
el_mask, dtype=np.float32 if opts.dtype == "fp32" else np.float16)
values = np.random.rand(n_els)
masked_rhs[np.nonzero(el_mask)] = values
if opts.block_size == 1:
triplets = sparse.triplets_from_dense(masked_rhs)
else:
triplets = sparse.triplets_from_dense(block_mask)
triplets = sparse.Triplets(
triplets.row_indices, triplets.col_indices,
sparse.blocks_at_indices(triplets.row_indices,
triplets.col_indices, opts.block_size,
masked_rhs))
fc = layers.SparseFcLayer.from_triplets(opts.output_size,
[opts.batchsize, opts.input_size],
*triplets,
matmul_options=matmul_opts,
name="fc_None",
dtype=dtype,
use_bias=False,
relu=False,
pooling_type='NONE')
fc_pool = layers.SparseFcLayer.from_triplets(
opts.output_size, [opts.batchsize, opts.input_size],
*triplets,
matmul_options=matmul_opts,
name="fc_" + opts.pooling_type,
dtype=dtype,
use_bias=False,
relu=False,
pooling_type=opts.pooling_type)
return fc, fc_pool
def test_representation_round_trip_blocks(self):
from ipu_sparse_ops import sparse
for bs in [4, 8, 16]:
# Create a mask that describes the non-zero block structure:
block_mask = np.array([[1, 1, 0], [0, 1, 0], [1, 0, 0], [0, 0, 1]])
n_blocks = np.count_nonzero(block_mask)
# From that produce an element-wise mask using a Kronecker product:
mask = np.kron(block_mask, np.ones(shape=[bs, bs])).astype(int)
n_els = np.count_nonzero(mask)
# Make a dense matrix from the element-wise mask and fill with random values:
dense = np.zeros_like(mask, dtype=np.float32)
values = np.random.rand(n_els)
dense[np.nonzero(mask)] = values
# Make the spec for the sparse matmul:
opts = {"metaInfoBucketOversizeProportion": 1}
spec = sparse.matmul_spec_from_max(dense.shape[1],
[2, dense.shape[0]],
max_non_zeros=n_blocks,
block_size=bs,
dtype=tf.float32)
# Make triplets indices from the block mask:
t = sparse.triplets_from_dense(block_mask)
# Then fill in triplet's values by extracting the blocks
# from the dense matrix (this can't be done by reshaping):
t_block = sparse.Triplets(
t.row_indices, t.col_indices,
sparse.blocks_at_indices(t.row_indices, t.col_indices, bs,
dense))
# Convert to on device representation and back and check the
# result is the dense matrix we sytarted with:
r = sparse.representation_from_triplets(spec, *t_block, opts)
t_rt = sparse.triplets_from_representation(spec, r, opts)
dense_rt = sparse.dense_from_triplets(spec, *t_rt)
assert_equal(dense, dense_rt)
# Check triplets from dense returns original triplets:
td = sparse.triplets_from_dense(dense_rt, bs)
assert_equal(t_block.row_indices, td.row_indices)
assert_equal(t_block.col_indices, td.col_indices)
assert_equal(t_block.values, td.values)
def make_fc_layer_and_test_inputs(args):
input_size = args.input_size
output_size = args.output_size
batch_size = args.batch_size
weights_type = tf.float16 if args.data_type == 'fp16' else tf.float32
matmul_opts = {"metaInfoBucketOversizeProportion": args.meta_info_oversize}
if args.pattern == 'fixed':
in_blocks = input_size // args.block_size
out_blocks = output_size // args.block_size
identity_size = max(in_blocks, out_blocks)
block_mask = np.identity(identity_size)[0:in_blocks, 0:out_blocks]
block_mask[1, 3] = 1
block_mask[0, 7] = 1
n_blocks = np.count_nonzero(block_mask)
el_mask = sparse.block_mask_to_element(block_mask, args.block_size)
n_els = np.count_nonzero(el_mask)
masked_rhs = np.zeros_like(el_mask, dtype=np.float32)
values = np.random.rand(n_els)
masked_rhs[np.nonzero(el_mask)] = values
if args.block_size == 1:
triplets = sparse.triplets_from_dense(masked_rhs)
else:
triplets = sparse.triplets_from_dense(block_mask)
triplets = sparse.Triplets(
triplets.row_indices, triplets.col_indices,
sparse.blocks_at_indices(triplets.row_indices,
triplets.col_indices, args.block_size,
masked_rhs))
fc = layers.SparseFcLayer.from_triplets(
args.output_size, [args.batch_size, args.input_size],
*triplets,
matmul_options=matmul_opts,
name='sparse_fc_from_triplets',
dtype=weights_type,
use_bias=False,
relu=False,
pooling_type=args.pooling_type)
elif args.pattern == 'random_sign_ones':
indices_random_gen = np.random.default_rng(seed=random_seed)
fc = layers.SparseFcLayer.from_random_generator(
args.output_size, [args.batch_size, args.input_size],
args.density,
args.block_size,
random_sign_ones_generator,
indices_random_gen,
matmul_options=matmul_opts,
name='sparse_fc_from_random_sign_ones',
use_bias=False,
relu=False,
pooling_type=args.pooling_type)
masked_rhs = sparse.dense_from_triplets(fc.weights.spec,
*fc.weights.triplets)
elif args.pattern == "random_orthogonal":
if args.input_size != args.output_size:
raise ValueError(
"random_orthogonal pattern requires square matrix")
matrix, max_non_zeros = sparse.gen_sparse_rand_orthog_mat(
args.output_size, args.density, args.block_size)
triplets = sparse.triplets_from_dense(matrix, args.block_size)
fc = layers.SparseFcLayer.from_triplets(
args.output_size, [args.batch_size, args.input_size],
*triplets,
matmul_options=matmul_opts,
name='sparse_fc_random_orthogonal',
dtype=weights_type,
use_bias=False,
relu=False,
pooling_type=args.pooling_type)
masked_rhs = sparse.dense_from_triplets(fc.weights.spec,
*fc.weights.triplets)
else:
random_gen = np.random.default_rng(seed=random_seed)
indices_random_gen = np.random.default_rng(seed=random_seed)
fc = layers.SparseFcLayer.from_random_generator(
args.output_size, [args.batch_size, args.input_size],
args.density,
args.block_size,
random_gen.standard_normal,
indices_random_gen,
matmul_options=matmul_opts,
name='sparse_fc_from_random',
dtype=weights_type,
use_bias=False,
relu=False,
pooling_type=args.pooling_type)
masked_rhs = fc.weights.extract_dense()
return fc, masked_rhs.astype(weights_type.as_numpy_dtype())
})
# Check all the results:
# Convert the sparse gradient metainfo back to triplets and then use those row and col indices
# to index the dense reference weight gradient:
sparse_data = sparse.SparseRepresentation(fc.weights.get_metainfo(),
sparse_weight_grad[0])
triplets = sparse.triplets_from_representation(fc.weights.spec,
sparse_data,
fc.weights.matmul_options)
if args.block_size == 1:
reference_grad_nzvalues = sparse.values_at_indices(
triplets[0], triplets[1], reference_weight_grad)
else:
reference_grad_nzvalues = sparse.blocks_at_indices(
triplets[0], triplets[1], args.block_size, reference_weight_grad)
# Convert the dense reference weight gradient to a sparse one using the same mask
# that we used for the weights so we can compare the nzvalues against the sparse grad:
dense_data = sparse.representation_from_triplets(fc.weights.spec,
triplets[0], triplets[1],
reference_grad_nzvalues,
fc.weights.matmul_options)
# Set tolerances appropriately as numpy is set for doubles by default:
if args.pattern == 'random_sign_ones':
rtol = 0
atol = 0
elif args.data_type == 'fp16':
rtol = 1e-03
atol = 1e-02
elif args.pattern == 'random_orthogonal':