def testPerLayerBlockSparsity(self):
param_list = [
"block_dims_map=[layer1/weights:1x1,layer2/weights:1x2]",
"block_pooling_function=AVG", "threshold_decay=0.0"
]
test_spec = ",".join(param_list)
pruning_hparams = pruning.get_pruning_hparams().parse(test_spec)
with tf.variable_scope("layer1"):
w1 = tf.constant([[-0.1, 0.1], [-0.2, 0.2]], name="weights")
pruning.apply_mask(w1)
with tf.variable_scope("layer2"):
w2 = tf.constant([[0.1, 0.1, 0.3, 0.3], [0.2, 0.2, 0.4, 0.4]],
name="weights")
pruning.apply_mask(w2)
sparsity = tf.Variable(0.5, name="sparsity")
p = pruning.Pruning(pruning_hparams, sparsity=sparsity)
mask_update_op = p.mask_update_op()
with self.cached_session() as session:
tf.global_variables_initializer().run()
session.run(mask_update_op)
mask1_eval = session.run(pruning.get_masks()[0])
mask2_eval = session.run(pruning.get_masks()[1])
self.assertAllEqual(session.run(pruning.get_weight_sparsity()),
[0.5, 0.5])
self.assertAllEqual(mask1_eval, [[0.0, 0.0], [1., 1.]])
self.assertAllEqual(mask2_eval, [[0, 0, 1., 1.], [0, 0, 1., 1.]])
def testConditionalMaskUpdate(self):
param_list = [
"pruning_frequency=2", "begin_pruning_step=1",
"end_pruning_step=6", "nbins=100"
]
test_spec = ",".join(param_list)
pruning_hparams = pruning.get_pruning_hparams().parse(test_spec)
weights = tf.Variable(tf.linspace(1.0, 100.0, 100), name="weights")
masked_weights = pruning.apply_mask(weights)
sparsity = tf.Variable(0.00, name="sparsity")
# Set up pruning
p = pruning.Pruning(pruning_hparams, sparsity=sparsity)
p._spec.threshold_decay = 0.0
mask_update_op = p.conditional_mask_update_op()
sparsity_val = tf.linspace(0.0, 0.9, 10)
increment_global_step = tf.assign_add(self.global_step, 1)
non_zero_count = []
with self.cached_session() as session:
tf.global_variables_initializer().run()
for i in range(10):
session.run(tf.assign(sparsity, sparsity_val[i]))
session.run(mask_update_op)
session.run(increment_global_step)
non_zero_count.append(np.count_nonzero(masked_weights.eval()))
# Weights pruned at steps 0,2,4,and,6
expected_non_zero_count = [100, 100, 80, 80, 60, 60, 40, 40, 40, 40]
self.assertAllEqual(expected_non_zero_count, non_zero_count)
def testWeightSpecificSparsity(self):
param_list = [
"begin_pruning_step=1", "pruning_frequency=1",
"end_pruning_step=100", "target_sparsity=0.5",
"weight_sparsity_map=[layer1:0.6,layer2/weights:0.75,.*kernel:0.6]",
"threshold_decay=0.0"
]
test_spec = ",".join(param_list)
pruning_hparams = pruning.get_pruning_hparams().parse(test_spec)
with tf.variable_scope("layer1"):
w1 = tf.Variable(tf.linspace(1.0, 100.0, 100), name="weights")
_ = pruning.apply_mask(w1)
with tf.variable_scope("layer2"):
w2 = tf.Variable(tf.linspace(1.0, 100.0, 100), name="weights")
_ = pruning.apply_mask(w2)
with tf.variable_scope("layer3"):
w3 = tf.Variable(tf.linspace(1.0, 100.0, 100), name="kernel")
_ = pruning.apply_mask(w3)
p = pruning.Pruning(pruning_hparams)
mask_update_op = p.conditional_mask_update_op()
increment_global_step = tf.assign_add(self.global_step, 1)
with self.cached_session() as session:
tf.global_variables_initializer().run()
for _ in range(110):
session.run(mask_update_op)
session.run(increment_global_step)
self.assertAllClose(session.run(pruning.get_weight_sparsity()),
[0.6, 0.75, 0.6])
def setUp(self):
super(PruningHParamsTest, self).setUp()
# Add global step variable to the graph
self.global_step = tf.train.get_or_create_global_step()
# Add sparsity
self.sparsity = tf.Variable(0.5, name="sparsity")
# Parse hparams
self.pruning_hparams = pruning.get_pruning_hparams().parse(
self.TEST_HPARAMS)
def _blockMasking(self, hparams, weights, expected_mask):
threshold = tf.Variable(0.0, name="threshold")
sparsity = tf.Variable(0.5, name="sparsity")
test_spec = ",".join(hparams)
pruning_hparams = pruning.get_pruning_hparams().parse(test_spec)
# Set up pruning
p = pruning.Pruning(pruning_hparams, sparsity=sparsity)
with self.cached_session():
tf.global_variables_initializer().run()
_, new_mask = p._maybe_update_block_mask(weights, threshold)
# Check if the mask is the same size as the weights
self.assertAllEqual(new_mask.get_shape(), weights.get_shape())
mask_val = new_mask.eval()
self.assertAllEqual(mask_val, expected_mask)