def test_prune_by_bbb_from_scratch_is_correct(self):
hidden_layers = [256, 128]
sampling_mode_ph = array_ops.placeholder(dtypes.string, [])
get_bbb_variable_fn = bayes_by_backprop.bayes_by_backprop_getter(
posterior_builder=bayes_by_backprop.diagonal_gaussian_posterior_builder,
prior_builder=bayes_by_backprop.adaptive_gaussian_prior_builder,
kl_builder=bayes_by_backprop.stochastic_kl_builder,
sampling_mode_tensor=sampling_mode_ph)
# create the bayes network
bbb_scope = 'net'
inputs_ph = array_ops.placeholder(dtypes.float32, [None, 784])
with variable_scope.variable_scope(bbb_scope, custom_getter=get_bbb_variable_fn) as vs:
logits_bbb = mlp(inputs_ph, hidden_layers)
# create the pruning op
metadata = bayes_by_backprop.get_variable_metadata()
total_variables = sum([np.prod(meta.raw_variable_shape) for meta in metadata])
percentage_ph = array_ops.placeholder(dtypes.float32, [])
pruned_vars_op = bbb.prune_by_bbb(metadata, percentage_ph)
# create the template network
template_scope = 'template'
with variable_scope.variable_scope(template_scope) as vs:
logits = mlp(inputs_ph, hidden_layers)
# retreve the variables so we can prune them in-place
template_variables = variables.trainable_variables(scope=template_scope)
# find the variables from 'net' that correspond to 'template'
assign_pruned_vars_op = bbb.assign_pruned_by_bbb_to_template(
metadata, pruned_vars_op, template_variables,
from_scope=bbb_scope, to_scope=template_scope)
with self.test_session() as sess:
sess.run(variables.global_variables_initializer())
for percentage in np.arange(0., 1., .05):
# ordinary pruning without the test-case looks like this:
# >>> sess.run(assign_pruned_vars_op, feed_dict={percentage_ph: percentage})
test_variables, _ = sess.run(
(pruned_vars_op, assign_pruned_vars_op),
feed_dict={
percentage_ph: percentage,
sampling_mode_ph: bayes_by_backprop.EstimatorModes.mean,
})
nonzero = 0
for test_variable, template_variable in zip(test_variables, template_variables):
nonzero += np.count_nonzero(test_variable)
self.assertAllClose(test_variable, sess.run(template_variable), atol=1e-8)
self.assertAlmostEqual(int(nonzero), int(total_variables * (1 - percentage)), delta=1)
def test_prune_by_bbb_from_scratch_to_sparse_ops(self):
hidden_layers = [256, 128]
sampling_mode_ph = array_ops.placeholder(dtypes.string, [])
get_bbb_variable_fn = bayes_by_backprop.bayes_by_backprop_getter(
posterior_builder=bayes_by_backprop.diagonal_gaussian_posterior_builder,
prior_builder=bayes_by_backprop.adaptive_gaussian_prior_builder,
kl_builder=bayes_by_backprop.stochastic_kl_builder,
sampling_mode_tensor=sampling_mode_ph)
# create the bayes network
bbb_scope = 'net'
inputs_ph = array_ops.placeholder(dtypes.float32, [None, 784])
with variable_scope.variable_scope(bbb_scope, custom_getter=get_bbb_variable_fn) as vs:
logits_bbb = mlp(inputs_ph, hidden_layers)
# create the optimal pruning op
metadata = bayes_by_backprop.get_variable_metadata()
total_variables = sum([np.prod(meta.raw_variable_shape) for meta in metadata])
percentage_ph = array_ops.placeholder(dtypes.float32, [])
pruned_vars_op = bbb.prune_by_bbb(metadata, percentage_ph)
# create the template network
template_scope = 'template'
with variable_scope.variable_scope(template_scope) as vs:
sparse_logits = mlp(inputs_ph, hidden_layers)
# retreve the variables so we can prune them in-place
template_variables = variables.trainable_variables(scope=template_scope)
# find the variables from 'net' that correspond to 'template'
assign_pruned_vars_op = bbb.assign_pruned_by_bbb_to_template(
metadata, pruned_vars_op, template_variables,
from_scope=bbb_scope, to_scope=template_scope)
with self.test_session() as sess:
test_inputs_ones = np.ones([1, 784])
sess.run(variables.global_variables_initializer())
for percentage in np.arange(0., 1.01, .01):
sess.run(
assign_pruned_vars_op, feed_dict={
percentage_ph: percentage,
sampling_mode_ph: bayes_by_backprop.EstimatorModes.mean,
})
sess.run(sparse_logits, feed_dict={inputs_ph: test_inputs_ones})
def test_mean_mode_is_deterministic_and_correct(self):
softplus_of_three = softplus(3.0)
bbb_getter = bbb.bayes_by_backprop_getter(
posterior_builder=test_diag_gaussian_builder_builder(10.9, 3.0),
prior_builder=bbb.fixed_gaussian_prior_builder,
sampling_mode_tensor=tf.constant(bbb.EstimatorModes.mean))
with tf.variable_scope("my_scope", custom_getter=bbb_getter):
my_variable = tf.get_variable("v", shape=[2], dtype=tf.float32)
init_op = tf.group(tf.global_variables_initializer(),
tf.local_variables_initializer())
with self.test_session() as sess:
sess.run(init_op)
variable_value_one = sess.run(my_variable)
variable_value_two = sess.run(my_variable)
variable_value_three = sess.run(my_variable)
self.assertAllClose(variable_value_one,
np.zeros(shape=[2]) + 10.9,
atol=1e-5)
self.assertAllClose(variable_value_two,
np.zeros(shape=[2]) + 10.9,
atol=1e-5)
self.assertAllClose(variable_value_three,
np.zeros(shape=[2]) + 10.9,
atol=1e-5)
variable_metadata = bbb.get_variable_metadata()
self.assertTrue(len(variable_metadata) == 1)
q_dist_sigma = variable_metadata[0].posterior.scale
with self.test_session() as sess:
sigma_res = sess.run(q_dist_sigma)
self.assertAllClose(sigma_res,
np.zeros(shape=[2]) + softplus_of_three,
atol=1e-5)
def test_sample_mode_is_stochastic_and_can_be_switched(self):
use_mean = tf.constant(bbb.EstimatorModes.mean)
use_sample = tf.constant(bbb.EstimatorModes.sample)
sampling_mode = tf.get_variable(
"bbb_sampling_mode",
initializer=tf.constant_initializer(bbb.EstimatorModes.sample),
dtype=tf.string,
shape=(),
trainable=False)
set_to_mean_mode = tf.assign(sampling_mode, use_mean)
set_to_sample_mode = tf.assign(sampling_mode, use_sample)
softplus_of_twenty = softplus(20.0)
bbb_getter = bbb.bayes_by_backprop_getter(
posterior_builder=test_diag_gaussian_builder_builder(10.9, 20.0),
prior_builder=bbb.fixed_gaussian_prior_builder,
sampling_mode_tensor=sampling_mode)
with tf.variable_scope("my_scope", custom_getter=bbb_getter):
my_variable = tf.get_variable("v", shape=[10, 3], dtype=tf.float32)
# Check that the distribution has the right parameters.
variable_metadata = bbb.get_variable_metadata()
self.assertTrue(len(variable_metadata) == 1)
q_dist_mean = variable_metadata[0].posterior.loc
q_dist_sigma = variable_metadata[0].posterior.scale
init_op = tf.group(tf.global_variables_initializer(),
tf.local_variables_initializer())
with self.test_session() as sess:
sess.run(init_op)
mean_res, sigma_res = sess.run([q_dist_mean, q_dist_sigma])
variable_value_one = sess.run(my_variable)
variable_value_two = sess.run(my_variable)
self.assertAllClose(mean_res, np.zeros(shape=[10, 3])+10.9)
self.assertAllClose(sigma_res, np.zeros(shape=[10, 3]) + softplus_of_twenty)
actual_distance = np.sqrt(
np.sum(np.square(variable_value_one - variable_value_two)))
expected_distance_minimum = 5
self.assertGreater(actual_distance, expected_distance_minimum)
# Now the value should be deterministic again.
with self.test_session() as sess:
sess.run(set_to_mean_mode)
variable_value_three = sess.run(my_variable)
variable_value_four = sess.run(my_variable)
variable_value_five = sess.run(my_variable)
self.assertAllClose(variable_value_three,
np.zeros(shape=[10, 3]) + 10.9,
atol=1e-5)
self.assertAllClose(variable_value_four,
np.zeros(shape=[10, 3]) + 10.9,
atol=1e-5)
self.assertAllClose(variable_value_five,
np.zeros(shape=[10, 3]) + 10.9,
atol=1e-5)
# Now it should be stochastic again.
with self.test_session() as sess:
sess.run(set_to_sample_mode)
variable_value_six = sess.run(my_variable)
variable_value_seven = sess.run(my_variable)
actual_new_distance = np.sqrt(
np.sum(np.square(variable_value_six - variable_value_seven)))
self.assertGreater(actual_new_distance, expected_distance_minimum)