def test_adaptive_stage_using_state_update_tensors(self):
"""Tests adaptive encoding stage with state update tensors."""
encoder = core_encoder.EncoderComposer(
test_utils.AdaptiveNormalizeEncodingStage()).add_parent(
test_utils.PlusOneEncodingStage(), P1_VALS).make()
x = tf.constant(1.0)
state = encoder.initial_state()
for _ in range(1, 5):
initial_state = state
encode_params, decode_params = encoder.get_params(state)
encoded_x, state_update_tensors, input_shapes = encoder.encode(
x, encode_params)
decoded_x = encoder.decode(encoded_x, decode_params, input_shapes)
state = encoder.update_state(initial_state, state_update_tensors)
data = self.evaluate(
test_utils.TestData(x, encoded_x, decoded_x, initial_state,
state_update_tensors, state))
self.assertAllClose(data.x, data.decoded_x)
self.assertLessEqual(
data.initial_state[CHILDREN][P1_VALS][STATE][AN_FACTOR_STATE],
1.0)
self.assertEqual(
data.state_update_tensors[CHILDREN][P1_VALS][TENSORS]
[AN_NORM_UPDATE], 2.0)
self.assertLessEqual(data.encoded_x[P1_VALS][AN_VALS], 2.0)
def test_larger_num_iters_improves_accuracy(self):
# If last_iter_clip = True, this potentially computes lossy representation.
# Set delta large, to measure the effect of changing num_iters on accuracy.
x = np.random.randn(3, 12).astype(np.float32)
errors = []
seed = tf.constant([1, 2], tf.int64)
for num_iters in [2, 3, 4, 5]:
stage = kashin.KashinHadamardEncodingStage(num_iters=num_iters,
eta=0.9,
delta=100.0,
last_iter_clip=True)
encode_params, decode_params = stage.get_params()
# To keep the experiment consistent, we always need to use fixed seed.
encode_params[
kashin.KashinHadamardEncodingStage.SEED_PARAMS_KEY] = seed
decode_params[
kashin.KashinHadamardEncodingStage.SEED_PARAMS_KEY] = seed
encoded_x, decoded_x = self.encode_decode_x(
stage, x, encode_params, decode_params)
test_data = test_utils.TestData(x, encoded_x, decoded_x)
test_data = self.evaluate_test_data(test_data)
errors.append(np.linalg.norm(test_data.x - test_data.decoded_x))
for e1, e2 in zip(errors[:-1], errors[1:]):
# The incurred error with less iterations should be greater.
self.assertGreater(e1, e2)
def test_input_types(self, x_dtype):
# Tests different input dtypes.
x = tf.constant([1.0, 0.1, 0.01, 0.001, 0.0001], dtype=x_dtype)
threshold = 0.05
stage = misc.SplitBySmallValueEncodingStage(threshold=threshold)
encode_params, decode_params = stage.get_params()
encoded_x, decoded_x = self.encode_decode_x(stage, x, encode_params,
decode_params)
test_data = test_utils.TestData(x, encoded_x, decoded_x)
test_data = self.evaluate_test_data(test_data)
self._assert_is_integer(test_data.encoded_x[
misc.SplitBySmallValueEncodingStage.ENCODED_INDICES_KEY])
# The numpy arrays must have the same dtype as the arrays from test_data.
expected_encoded_values = np.array([1.0, 0.1],
dtype=x.dtype.as_numpy_dtype)
expected_encoded_indices = np.array([[0], [1]], dtype=np.int32)
expected_decoded_x = np.array([1.0, 0.1, 0., 0., 0.],
dtype=x_dtype.as_numpy_dtype)
self.assertAllEqual(test_data.encoded_x[stage.ENCODED_VALUES_KEY],
expected_encoded_values)
self.assertAllEqual(test_data.encoded_x[stage.ENCODED_INDICES_KEY],
expected_encoded_indices)
self.assertAllEqual(test_data.decoded_x, expected_decoded_x)
def test_input_types(self, x_dtype):
# Tests combinations of input dtypes.
stage = quantization.PRNGUniformQuantizationEncodingStage(bits=8)
x = tf.random.normal([50], dtype=x_dtype)
encode_params, decode_params = stage.get_params()
encoded_x, decoded_x = self.encode_decode_x(stage, x, encode_params,
decode_params)
test_data = test_utils.TestData(x, encoded_x, decoded_x)
test_data = self.evaluate_test_data(test_data)
def test_eta_delta_take_tf_values(self):
x = self.default_input()
stage = kashin.KashinHadamardEncodingStage(eta=tf.constant(0.9),
delta=tf.constant(1.0))
encode_params, decode_params = stage.get_params()
encoded_x, decoded_x = self.encode_decode_x(stage, x, encode_params,
decode_params)
test_data = test_utils.TestData(x, encoded_x, decoded_x)
self.generic_asserts(test_data, stage)
self.common_asserts_for_test_data(test_data)
def test_input_types(self, x_dtype, eta_dtype, delta_dtype):
stage = kashin.KashinHadamardEncodingStage(
eta=tf.constant(0.9, eta_dtype),
delta=tf.constant(1.0, delta_dtype))
x = tf.random.normal([3, 12], dtype=x_dtype)
encode_params, decode_params = stage.get_params()
encoded_x, decoded_x = self.encode_decode_x(stage, x, encode_params,
decode_params)
test_data = test_utils.TestData(x, encoded_x, decoded_x)
test_data = self.evaluate_test_data(test_data)
self.assertAllEqual(test_data.x.shape, test_data.decoded_x.shape)
def test_encoding_differs_given_different_seed(self):
"""Tests that encoded_x is different in different evaluations."""
x = tf.constant(self.evaluate(self.default_input()))
stage = self.default_encoding_stage()
encode_params, decode_params = stage.get_params()
encoded_x, decoded_x = self.encode_decode_x(stage, x, encode_params,
decode_params)
test_data_1 = self.evaluate_test_data(
test_utils.TestData(x, encoded_x, decoded_x))
test_data_2 = self.evaluate_test_data(
test_utils.TestData(x, encoded_x, decoded_x))
# The decoded values should be the sam, but the encoded values not.
self.assertAllClose(test_data_1.decoded_x,
test_data_2.decoded_x,
rtol=test_utils.DEFAULT_RTOL,
atol=test_utils.DEFAULT_ATOL)
self.assertNotAllClose(test_data_1.encoded_x[self._ENCODED_VALUES_KEY],
test_data_2.encoded_x[self._ENCODED_VALUES_KEY],
rtol=test_utils.DEFAULT_RTOL,
atol=test_utils.DEFAULT_ATOL)
def test_basic_encode_decode_tf_constructor_parameters(self):
"""Tests the core funcionality with `tf.Variable` constructor parameters."""
a_var = tf.get_variable('a_var', initializer=self._DEFAULT_A)
b_var = tf.get_variable('b_var', initializer=self._DEFAULT_B)
stage = test_utils.SimpleLinearEncodingStage(a_var, b_var)
with self.cached_session() as sess:
sess.run(tf.global_variables_initializer())
x = self.default_input()
encode_params, decode_params = stage.get_params()
encoded_x, decoded_x = self.encode_decode_x(stage, x, encode_params,
decode_params)
test_data = self.evaluate_test_data(
test_utils.TestData(x, encoded_x, decoded_x))
self.common_asserts_for_test_data(test_data)
# Change the variables and verify the behavior of stage changes.
self.evaluate([tf.assign(a_var, 5.0), tf.assign(b_var, 6.0)])
test_data = self.evaluate_test_data(
test_utils.TestData(x, encoded_x, decoded_x))
self.assertAllClose(test_data.x * 5.0 + 6.0,
test_data.encoded_x[self._ENCODED_VALUES_KEY])
def test_input_types(self, x_dtype, clip_norm_dtype):
# Tests combinations of input dtypes.
stage = clipping.ClipByNormEncodingStage(
tf.constant(1.0, clip_norm_dtype))
x = tf.constant([1.0, 1.0, 1.0, 1.0], dtype=x_dtype)
encode_params, decode_params = stage.get_params()
encoded_x, decoded_x = self.encode_decode_x(stage, x, encode_params,
decode_params)
test_data = test_utils.TestData(x, encoded_x, decoded_x)
test_data = self.evaluate_test_data(test_data)
self.assertAllEqual([1.0, 1.0, 1.0, 1.0], test_data.x)
# The decoded values should have norm 1.
self.assertAllClose([0.5, 0.5, 0.5, 0.5], test_data.decoded_x)
def test_empty_input_static(self):
# Tests that the encoding works when the input shape is [0].
x = []
x = tf.convert_to_tensor(x, dtype=tf.int32)
assert x.shape.as_list() == [0]
stage = self.default_encoding_stage()
encode_params, decode_params = stage.get_params()
encoded_x, decoded_x = self.encode_decode_x(stage, x, encode_params,
decode_params)
test_data = self.evaluate_test_data(
test_utils.TestData(x, encoded_x, decoded_x))
self.common_asserts_for_test_data(test_data)
def test_input_types(self, x_dtype, min_max_dtype):
# Tests combinations of input dtypes.
stage = stages_impl.UniformQuantizationEncodingStage(
bits=8, min_max=tf.constant([-1.0, 1.0], min_max_dtype))
x = tf.random.normal([50], dtype=x_dtype)
encode_params, decode_params = stage.get_params()
encoded_x, decoded_x = self.encode_decode_x(stage, x, encode_params,
decode_params)
test_data = test_utils.TestData(x, encoded_x, decoded_x)
test_data = self.evaluate_test_data(test_data)
self.assertLess(np.amin(test_data.x), -1.0)
self.assertGreater(np.amax(test_data.x), 1.0)
self.assertAllGreaterEqual(test_data.decoded_x, -1.0)
self.assertAllLessEqual(test_data.decoded_x, 1.0)
def test_adaptive_stage(self):
"""Tests composition of two adaptive encoding stages."""
encoder = core_encoder.EncoderComposer(
test_utils.PlusOneOverNEncodingStage()).add_parent(
test_utils.PlusOneOverNEncodingStage(), PN_VALS).make()
x = tf.constant(1.0)
state = encoder.initial_state()
for i in range(1, 5):
initial_state = state
encode_params, decode_params = encoder.get_params(state)
encoded_x, state_update_tensors, input_shapes = encoder.encode(
x, encode_params)
decoded_x = encoder.decode(encoded_x, decode_params, input_shapes)
state = encoder.update_state(initial_state, state_update_tensors)
data = self.evaluate(
test_utils.TestData(x, encoded_x, decoded_x, initial_state,
state_update_tensors, state))
expected_initial_state = {
STATE: {
PN_ITER_STATE: i
},
CHILDREN: {
PN_VALS: {
STATE: {
PN_ITER_STATE: i
},
CHILDREN: {}
}
}
}
expected_state_update_tensors = {
TENSORS: {},
CHILDREN: {
PN_VALS: {
TENSORS: {},
CHILDREN: {}
}
}
}
self.assertAllClose(data.x, data.decoded_x)
self.assertAllEqual(expected_initial_state, data.initial_state)
self.assertDictEqual(expected_state_update_tensors,
data.state_update_tensors)
self.assertAllClose(data.x + 2 * 1 / i,
data.encoded_x[PN_VALS][PN_VALS])
def test_input_types(self, x_dtype, clip_value_min_dtype,
clip_value_max_dtype):
# Tests combinations of input dtypes.
stage = clipping.ClipByValueEncodingStage(
tf.constant(-1.0, clip_value_min_dtype),
tf.constant(1.0, clip_value_max_dtype))
x = tf.constant([-2.0, -1.0, 0.0, 1.0, 2.0], dtype=x_dtype)
encode_params, decode_params = stage.get_params()
encoded_x, decoded_x = self.encode_decode_x(stage, x, encode_params,
decode_params)
test_data = test_utils.TestData(x, encoded_x, decoded_x)
test_data = self.evaluate_test_data(test_data)
self.common_asserts_for_test_data(test_data)
self.assertAllEqual([-2.0, -1.0, 0.0, 1.0, 2.0], test_data.x)
self.assertAllClose([-1.0, -1.0, 0.0, 1.0, 1.0], test_data.decoded_x)
def test_all_below_threshold_works(self):
# Tests that encoding does not blow up with all-below-threshold input. In
# this case, both of the encoded values will be empty arrays.
stage = misc.SplitBySmallValueEncodingStage(threshold=0.1)
x = tf.random.uniform([50], minval=-0.01, maxval=0.01)
encode_params, decode_params = stage.get_params()
encoded_x, decoded_x = self.encode_decode_x(stage, x, encode_params,
decode_params)
test_data = test_utils.TestData(x, encoded_x, decoded_x)
test_data = self.evaluate_test_data(test_data)
expected_encoded_indices = np.array([], dtype=np.int32).reshape([0, 1])
self.assertAllEqual(test_data.encoded_x[stage.ENCODED_VALUES_KEY], [])
self.assertAllEqual(test_data.encoded_x[stage.ENCODED_INDICES_KEY],
expected_encoded_indices)
self.assertAllEqual(test_data.decoded_x,
np.zeros([50], dtype=x.dtype.as_numpy_dtype))
def test_as_adaptive_encoding_stage(self):
"""Tests correctness of the wrapped encoding stage."""
a_var = tf.compat.v1.get_variable('a', initializer=2.0)
b_var = tf.compat.v1.get_variable('b', initializer=3.0)
stage = test_utils.SimpleLinearEncodingStage(a_var, b_var)
wrapped_stage = encoding_stage.as_adaptive_encoding_stage(stage)
self.assertIsInstance(wrapped_stage,
encoding_stage.AdaptiveEncodingStageInterface)
x = tf.constant(2.0)
state = wrapped_stage.initial_state()
encode_params, decode_params = wrapped_stage.get_params(state)
encoded_x, state_update_tensors = wrapped_stage.encode(
x, encode_params)
updated_state = wrapped_stage.update_state(state, state_update_tensors)
decoded_x = wrapped_stage.decode(encoded_x, decode_params)
# Test that the added state functionality is empty.
self.assertDictEqual({}, state)
self.assertDictEqual({}, state_update_tensors)
self.assertDictEqual({}, updated_state)
self.assertDictEqual({}, wrapped_stage.state_update_aggregation_modes)
# Test that __getattr__ retrieves attributes of the wrapped stage.
self.assertIsInstance(wrapped_stage._a, tf.Variable)
self.assertIs(wrapped_stage._a, a_var)
self.assertIsInstance(wrapped_stage._b, tf.Variable)
self.assertIs(wrapped_stage._b, b_var)
# Test the functionality remain unchanged.
self.assertEqual(stage.name, wrapped_stage.name)
self.assertEqual(stage.compressible_tensors_keys,
wrapped_stage.compressible_tensors_keys)
self.assertEqual(stage.commutes_with_sum,
wrapped_stage.commutes_with_sum)
self.assertEqual(stage.decode_needs_input_shape,
wrapped_stage.decode_needs_input_shape)
self.evaluate(tf.compat.v1.global_variables_initializer())
test_data = test_utils.TestData(
*self.evaluate([x, encoded_x, decoded_x]))
self.assertEqual(2.0, test_data.x)
self.assertEqual(
7.0, test_data.encoded_x[
test_utils.SimpleLinearEncodingStage.ENCODED_VALUES_KEY])
self.assertEqual(2.0, test_data.decoded_x)
def test_empty_input_dynamic(self):
# Tests that the encoding works when the input shape is [0], but not
# statically known.
y = tf.zeros((10, ))
indices = tf.compat.v2.where(tf.abs(y) > 1e-8)
x = tf.gather_nd(y, indices)
x = tf.cast(x, tf.int32) # Empty tensor.
assert x.shape.as_list() == [None]
stage = self.default_encoding_stage()
encode_params, decode_params = stage.get_params()
encoded_x, decoded_x = self.encode_decode_x(stage, x, encode_params,
decode_params)
test_data = self.evaluate_test_data(
test_utils.TestData(x, encoded_x, decoded_x))
assert test_data.x.shape == (0, )
assert test_data.encoded_x[stage.ENCODED_VALUES_KEY].shape == (0, )
assert test_data.decoded_x.shape == (0, )
def test_dynamic_input_shape(self):
# Tests that encoding works when the input shape is not statically known.
stage = quantization.PRNGUniformQuantizationEncodingStage(bits=8)
shape = [10, 5, 7]
prob = [0.5, 0.8, 0.6]
original_x = tf.random.uniform(shape, dtype=tf.float32)
rand = [tf.random.uniform([shape[i],]) for i in range(3)]
sample_indices = [
tf.reshape(tf.where(rand[i] < prob[i]), [-1]) for i in range(3)
]
x = tf.gather(original_x, sample_indices[0], axis=0)
x = tf.gather(x, sample_indices[1], axis=1)
x = tf.gather(x, sample_indices[2], axis=2)
encode_params, decode_params = stage.get_params()
encoded_x, decoded_x = self.encode_decode_x(stage, x, encode_params,
decode_params)
test_data = test_utils.TestData(x, encoded_x, decoded_x)
test_data = self.evaluate_test_data(test_data)
def test_quantization_empirically_unbiased(self):
# Tests that the quantization "seems" to be unbiased.
# Executing the encoding and decoding many times, the average error should
# be a lot larger than the error of average decoded value.
x = tf.constant(np.random.rand((50)).astype(np.float32))
stage = quantization.PRNGUniformQuantizationEncodingStage(bits=2)
encode_params, decode_params = stage.get_params()
encoded_x, decoded_x = self.encode_decode_x(stage, x, encode_params,
decode_params)
test_data = test_utils.TestData(x, encoded_x, decoded_x)
test_data_list = [self.evaluate_test_data(test_data) for _ in range(200)]
norm_errors = []
errors = []
for data in test_data_list:
norm_errors.append(np.linalg.norm(data.x - data.decoded_x))
errors.append(data.x - data.decoded_x)
mean_of_errors = np.mean(norm_errors)
error_of_mean = np.linalg.norm(np.mean(errors, axis=0))
self.assertGreater(mean_of_errors, error_of_mean * 10)
def test_approximately_unbiased_in_expectation(self):
"""Tests that average of encodings is more accurate than a single one."""
# Use a constant input value.
x = self.evaluate(self.default_input())
stage = self.default_encoding_stage()
encode_params, decode_params = stage.get_params()
encoded_x, decoded_x = self.encode_decode_x(stage, x, encode_params,
decode_params)
test_data = []
for _ in range(100):
test_data.append(
test_utils.TestData(
*self.evaluate_tf_py_list([x, encoded_x, decoded_x])))
# Check that the average error created by encoding is significantly larger
# than error of average of encodings. This is an simple (imperfect)
# empirical check that the encoding is unbiased.
mean_error = np.mean([np.linalg.norm(x - d.decoded_x) for d in test_data])
error_of_mean = np.linalg.norm(
x - np.mean([d.decoded_x for d in test_data], axis=0))
self.assertGreater(mean_error, error_of_mean * 5)
