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 from_triplets(cls,
hidden_size: int,
input_shape: List[int],
row_indices: List[int],
col_indices: List[int],
values: List[float],
matmul_options: Mapping[str, str],
name: str,
dtype: tf.DType = tf.float32,
use_bias: bool = False,
relu: bool = False,
disable_updating: bool = False,
pooling_type: str = 'NONE'):
"""
Utility factory function to build a 'SparseFcLayer' from a set of triplets (COO format).
E.g. as returned from 'ipu_sparse_ops.sparse.triplets_from_dense'
"""
block_size = sparse.block_size_from_list(values)
spec = sparse.matmul_spec_from_max(hidden_size, input_shape,
len(values), block_size, dtype,
pooling_type)
ns = tf.get_default_graph().get_name_scope()
qualified_name = ns + "/" + name if ns else name
logger.debug(
f"Creating random sparse FC {qualified_name} with spec: {spec}")
triplets = sparse.Triplets(row_indices, col_indices, values)
weights = SparseMatrix(spec,
matmul_options,
triplets,
name=qualified_name)
return cls(weights,
name,
use_bias,
relu,
disable_updating,
pooling_type=pooling_type)
def prune_and_grow(name,
triplets,
shape,
spec,
max_non_zeros,
slot_triplets,
prune_schedule,
prune_ratio: float,
grad_w: np.array,
grow_method='rigl',
random_gen=None,
ipu_pooling_type='NONE'):
"""
Performs the pruning and growing of the weights to update the sparsity pattern, for the current fc layer
:param name: A debug name
:param triplets: Current triplets to prune and grow
:param shape: Shape of the dense matrix
:param spec: Specs of the sparse matmul
:param max_non_zeros: Maximum number of non-zeros values
:slot_triplets: Triplets for the current slots
:param prune_schedule: a function which given max prune count returns the number of weights to update
:param prune_ratio: the maximum percentage of this layer's weights to update
:param grad_w: A numpy array containing the dense gradients for each sub-layer
:param grow_method: Method used to regrow the weights, either rigl or random selection
:param random_gen: Random number generator to be used if the grow_method is 'random'
"""
if isinstance(grad_w, list):
grad_w = grad_w[0]
if isinstance(grad_w, dict):
grad_w = [grad for grad in grad_w.values()][0]
# Compute the prune count using the provides schedule, the max number of zeros in the
# layer and the pruning ratio
prune_count = prune_schedule(
max_pruned=int(np.ceil(prune_ratio * max_non_zeros)))
if prune_count == 0:
logger.info("Nothing to prune according to prune schedule.")
return None
logger.info(
f"Triplet stats before prune and grow for {name}: {sparse.triplet_stats(*triplets)}"
)
if logger.level <= logging.DEBUG:
abs_nz_values = np.abs(triplets[2])
if len(abs_nz_values.shape) > 1:
abs_nz_values = abs_nz_values.sum(-1).sum(-1)
block_input_size = spec.input_size // spec.block_size
block_output_size = spec.output_size // spec.block_size
block_spec = sparse.MatmulSpec(
block_size=1,
input_size=block_input_size,
output_size=block_output_size,
num_groups=spec.num_groups,
batch_size=spec.batch_size,
data_type=spec.data_type,
max_non_zero_blocks=spec.max_non_zero_blocks,
pooling_type=spec.pooling_type)
dense_abs_weights = sparse.dense_from_triplets(
block_spec, triplets[0], triplets[1], abs_nz_values)
plot_and_log_matrix(name + "/abs_block_weights", dense_abs_weights)
# Prune bottom k weights
logging.debug(
f"Pruning and grow also applies to these slot vars: {slot_triplets.keys()}"
)
slot_values = {
name: triplet.values
for name, triplet in slot_triplets.items()
}
remaining_triplets, remaining_slot_values = prune_bottom_k_weights(
*triplets, slot_values, prune_count, name)
# regrow weights
logger.debug(
f"Regrowing non-zeros for layer {name} using '{grow_method}' method.")
if grow_method == 'rigl':
weights_shape = np.array(triplets[2]).shape
block_size = 1 if len(weights_shape) == 1 else weights_shape[-1]
# Grow back new indices using Rig-L: (https://arxiv.org/abs/1911.11134)
if (shape != grad_w.shape
and (ipu_pooling_type == "NONE" or block_size == 1)):
raise RuntimeError(
f"Dense weight gradient has unexpected shape.Expected {shape}, got {grad_w.shape}"
)
new_triplets = regrow_rigl(triplets, grad_w, zero_values_generator,
prune_count, ipu_pooling_type == "NONE",
name)
if grow_method == 'random':
# Random replacement strategy: add back random indices
# Gen some replacement random indices excluding all the existing
# ones then we will swap for the pruned ones:
new_triplets = sparse.random_triplets(
spec,
indices_initialiser_gen=random_gen,
value_generator=zero_values_generator,
excluded_indices=(triplets[0], triplets[1]),
count=prune_count)
grown_triplets, grown_slots = join_triplets(remaining_triplets,
new_triplets,
remaining_slot_values,
prune_count)
if len(grown_triplets[0]) != max_non_zeros:
raise ValueError(
f"Grown row count {len(grown_triplets[0])} does not match expected count {max_non_zeros}"
)
if len(grown_triplets[1]) != max_non_zeros:
raise ValueError(
f"Grown col count {len(grown_triplets[1])} does not match expected count {max_non_zeros}"
)
if len(grown_triplets[2]) != max_non_zeros:
raise ValueError(
f"Grown col count {len(grown_triplets[2])} does not match expected count {max_non_zeros}"
)
for grown_slot in grown_slots.values():
if len(grown_slot) != max_non_zeros:
raise ValueError(
f"Grown col count {len(grown_slot)} does not match expected count {max_non_zeros}"
)
grown_triplets = sparse.Triplets(grown_triplets[0], grown_triplets[1],
grown_triplets[2])
grown_slots = {
name: sparse.Triplets(grown_triplets[0], grown_triplets[1], grown_slot)
for name, grown_slot in grown_slots.items()
}
logger.info(
f"Triplet stats after prune and grow for {name}: {sparse.triplet_stats(*grown_triplets)}"
)
return {
'gt': grown_triplets,
'nt': new_triplets,
'rt': remaining_triplets,
'gs': grown_slots,
'name': name
}
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())
def prune_and_grow(name,
triplets,
shape,
spec,
max_non_zeros,
slot_triplets,
prune_schedule,
prune_ratio: float,
grad_w: np.array,
grow_method='rigl',
random_gen=None):
"""
Performs the pruning and growing of the weights to update the sparsity pattern, for the current fc layer
:param name: A debug name
:param triplets: Current triplets to prune and grow
:param shape: Shape of the dense matrix
:param spec: Specs of the sparse matmul
:param max_non_zeros: Maximum number of non-zeros values
:slot_triplets: Triplets for the current slots
:param prune_schedule: a function which given max prune count returns the number of weights to update
:param prune_ratio: the maximum percentage of this layer's weights to update
:param grad_w: A numpy array containing the dense gradients for each sub-layer
:param grow_method: Method used to regrow the weights, either rigl or random selection
:param random_gen: Random number generator to be used if the grow_method is 'random'
"""
if isinstance(grad_w, list):
grad_w = grad_w[0]
# Compute the prune count using the provides schedule, the max numbe rof zeros in the
# layer and the pruning ratio
prune_count = prune_schedule(
max_pruned=int(np.ceil(prune_ratio * max_non_zeros)))
# Prune bottom k weights
slot_values = {
name: triplet.values
for name, triplet in slot_triplets.items()
}
remaining_triplets, remaining_slot_values = prune_bottom_k_weights(
*triplets, slot_values, prune_count, name)
zero_values = zero_values_generator
# regrow weights
if grow_method == 'rigl':
# Grow back new indices using Rig-L: (https://arxiv.org/abs/1911.11134)
if shape != grad_w.shape:
raise RuntimeError(
f"Dense weight gradient has unexpected shape.Expected {shape}, got {grad_w.shape}"
)
new_triplets = regrow_rigl(triplets, grad_w, zero_values, prune_count,
name)
if grow_method == 'random':
# Random replacement strategy: add back random indices
# Gen some replacement random indices excluding all the existing
# ones then we will swap for the pruned ones:
new_triplets = sparse.random_triplets(
spec,
indices_initialiser_gen=random_gen,
value_generator=zero_values,
excluded_indices=(triplets[0], triplets[1]),
count=prune_count)
grown_triplets, grown_slots = join_triplets(remaining_triplets,
new_triplets,
remaining_slot_values,
prune_count)
if len(grown_triplets[0]) != max_non_zeros:
raise ValueError(
f"Grown row count {len(grown_triplets[0])} does not match expected count {max_non_zeros}"
)
if len(grown_triplets[1]) != max_non_zeros:
raise ValueError(
f"Grown col count {len(grown_triplets[1])} does not match expected count {max_non_zeros}"
)
if len(grown_triplets[2]) != max_non_zeros:
raise ValueError(
f"Grown col count {len(grown_triplets[2])} does not match expected count {max_non_zeros}"
)
for grown_slot in grown_slots.values():
if len(grown_slot) != max_non_zeros:
raise ValueError(
f"Grown col count {len(grown_slot)} does not match expected count {max_non_zeros}"
)
grown_triplets = sparse.Triplets(grown_triplets[0], grown_triplets[1],
grown_triplets[2])
grown_slots = {
name: [grown_triplets[0], grown_triplets[1], grown_slot]
for name, grown_slot in grown_slots.items()
}
return {
'gt': grown_triplets,
'nt': new_triplets,
'rt': remaining_triplets,
'gs': grown_slots
}