def _compare_expand_tensor_with_kronecker_product(self, tensor, block_dim):
with self.cached_session() as session:
compat.initialize_variables(self)
expanded_tensor = pruning_utils.expand_tensor(tensor, block_dim)
kronecker_product = pruning_utils.kronecker_product(
tensor, tf.ones(block_dim))
expanded_tensor_val, kronecker_product_val = session.run(
[expanded_tensor, kronecker_product])
self.assertAllEqual(expanded_tensor_val, kronecker_product_val)
def initialize(self):
self.global_step = tf.Variable(tf.zeros([], dtype=dtypes.int32),
dtype=dtypes.int32,
name="global_step")
def training_step_fn():
return self.global_step
self.training_step_fn = training_step_fn
compat.initialize_variables(self)
def testPrunesWithConstantSparsity(self):
sparsity = pruning_schedule.ConstantSparsity(0.5, 100, 200, 10)
step_100 = tf.Variable(100)
step_110 = tf.Variable(110)
step_200 = tf.Variable(200)
compat.initialize_variables(self)
self.assertAllClose(0.5, self.evaluate(sparsity(step_100))[1])
self.assertAllClose(0.5, self.evaluate(sparsity(step_110))[1])
self.assertAllClose(0.5, self.evaluate(sparsity(step_200))[1])
def testPrunesForeverIfEndStepIsNegativeOne(self):
sparsity = pruning_schedule.ConstantSparsity(0.5, 0, -1, 10)
step_10000 = tf.Variable(10000)
step_100000000 = tf.Variable(100000000)
compat.initialize_variables(self)
self.assertTrue(self.evaluate(sparsity(step_10000))[0])
self.assertTrue(self.evaluate(sparsity(step_100000000))[0])
self.assertAllClose(0.5, self.evaluate(sparsity(step_10000))[1])
self.assertAllClose(0.5, self.evaluate(sparsity(step_100000000))[1])
def _compare_pooling_methods(self, weights, pooling_kwargs):
with self.cached_session():
compat.initialize_variables(self)
pooled_weights_tf = tf.squeeze(
tf.nn.pool(
tf.reshape(
weights,
[1, weights.get_shape()[0],
weights.get_shape()[1], 1]), **pooling_kwargs))
pooled_weights_factorized_pool = pruning_utils.factorized_pool(
weights, **pooling_kwargs)
self.assertAllClose(self.evaluate(pooled_weights_tf),
self.evaluate(pooled_weights_factorized_pool))
def testPolynomialDecay_PrunesCorrectly(self):
sparsity = pruning_schedule.PolynomialDecay(0.2, 0.8, 100, 110, 3, 2)
step_100 = tf.Variable(100)
step_102 = tf.Variable(102)
step_105 = tf.Variable(105)
step_110 = tf.Variable(110)
compat.initialize_variables(self)
# These values were generated using tf.polynomial_decay with the same
# params in a colab to verify.
self.assertAllClose(0.2, self.evaluate(sparsity(step_100))[1])
self.assertAllClose(0.4928, self.evaluate(sparsity(step_102))[1])
self.assertAllClose(0.725, self.evaluate(sparsity(step_105))[1])
self.assertAllClose(0.8, self.evaluate(sparsity(step_110))[1])
def testOnlyPrunesAtValidFrequencySteps(self, schedule_type):
sparsity = self._construct_pruning_schedule(schedule_type, 100, 200,
10)
step_100 = tf.Variable(100)
step_109 = tf.Variable(109)
step_110 = tf.Variable(110)
step_111 = tf.Variable(111)
compat.initialize_variables(self)
self.assertFalse(self.evaluate(sparsity(step_109))[0])
self.assertFalse(self.evaluate(sparsity(step_111))[0])
self.assertTrue(self.evaluate(sparsity(step_100))[0])
self.assertTrue(self.evaluate(sparsity(step_110))[0])
def _GetMinMaxValues(self, quantize_fn, input_values, **kwds):
min_var = tf.Variable(0.0)
max_var = tf.Variable(0.0)
compat.initialize_variables(self)
for input_elem in input_values:
y = quantize_fn(input_elem,
min_var,
max_var,
is_training=True,
**kwds)
self.evaluate(y)
# Now check that the min_max_vars were, in fact, updated.
min_max_values = self.evaluate([min_var, max_var])
return min_max_values[0], min_max_values[1]
def _test_quantizer(self, quantizer):
inputs = tf.Variable(np.array([[-1.0, 0.5], [0.0, 1.0]]),
name='inputs',
dtype=tf.dtypes.float32)
min_var = tf.Variable(0.0)
max_var = tf.Variable(0.0)
kwargs = {'min_var': min_var, 'max_var': max_var}
quant_tensor = quantizer(inputs, step=0, training=True, **kwargs)
compat.initialize_variables(self)
results = self.evaluate(quant_tensor)
min_max_values = self.evaluate([min_var, max_var])
# TODO(pulkitb): Assert on expected values for testing.
# Since the underlying code is already tested in quant_ops_test.py, this
# just ensures the Quantizers code is wired properly.
print('Result: ', results)
print('min_var: ', min_max_values[0])
print('max_var: ', min_max_values[1])
def testPrunesOnlyInBeginEndStepRange(self, schedule_type):
sparsity = self._construct_pruning_schedule(schedule_type, 100, 200, 1)
# Before begin step
step_90 = tf.Variable(90)
step_99 = tf.Variable(99)
# In range
step_100 = tf.Variable(100)
step_110 = tf.Variable(110)
step_200 = tf.Variable(200)
# After end step
step_201 = tf.Variable(201)
step_210 = tf.Variable(210)
compat.initialize_variables(self)
self.assertFalse(self.evaluate(sparsity(step_90))[0])
self.assertFalse(self.evaluate(sparsity(step_99))[0])
self.assertTrue(self.evaluate(sparsity(step_100))[0])
self.assertTrue(self.evaluate(sparsity(step_110))[0])
self.assertTrue(self.evaluate(sparsity(step_200))[0])
self.assertFalse(self.evaluate(sparsity(step_201))[0])
self.assertFalse(self.evaluate(sparsity(step_210))[0])