def hadamard_quantization(bits):
"""Returns hadamard quanitzation `Encoder`.
The `Encoder` first reshapes the input to a rank-1 `Tensor`, and applies the
Hadamard transform (rotation). It then applies uniform quantization with the
extreme values being the minimum and maximum of the rotated vector being
encoded. Finally, the quantized values are bitpacked to an integer type.
The `Encoder` is a composition of the following encoding stages:
* `FlattenEncodingStage` - reshaping the input to a vector.
* `HadamardEncodingStage` - applying the Hadamard transform.
* `UniformQuantizationEncodingStage` - applying uniform quantization.
* `BitpackingEncodingStage` - bitpacking the result into integer values.
Args:
bits: Number of bits to quantize into.
Returns:
The hadamard quantization `Encoder`.
"""
return core_encoder.EncoderComposer(
stages_impl.BitpackingEncodingStage(bits)).add_parent(
stages_impl.UniformQuantizationEncodingStage(bits),
stages_impl.UniformQuantizationEncodingStage.ENCODED_VALUES_KEY
).add_parent(
stages_impl.HadamardEncodingStage(),
stages_impl.HadamardEncodingStage.ENCODED_VALUES_KEY).add_parent(
stages_impl.FlattenEncodingStage(),
stages_impl.FlattenEncodingStage.ENCODED_VALUES_KEY).make()
def uniform_quantization(bits):
"""Returns uniform quanitzation `Encoder`.
The `Encoder` first reshapes the input to a rank-1 `Tensor`, then applies
uniform quantization with the extreme values being the minimum and maximum of
the vector being encoded. Finally, the quantized values are bitpacked to an
integer type.
The `Encoder` is a composition of the following encoding stages:
* `FlattenEncodingStage`
* `UniformQuantizationEncodingStage`
* `BitpackingEncodingStage`
Args:
bits: Number of bits to quantize into.
Returns:
The quantization `Encoder`.
"""
return core_encoder.EncoderComposer(
stages_impl.BitpackingEncodingStage(bits)).add_parent(
stages_impl.UniformQuantizationEncodingStage(bits), stages_impl.
UniformQuantizationEncodingStage.ENCODED_VALUES_KEY).add_parent(
stages_impl.FlattenEncodingStage(),
stages_impl.FlattenEncodingStage.ENCODED_VALUES_KEY).make()
def test_all_zero_input_works(self):
# Tests that encoding does not blow up with all-zero input. With
# min_max=None, the derived min and max are identical, thus potential for
# division by zero.
stage = stages_impl.UniformQuantizationEncodingStage(bits=8, min_max=None)
test_data = self.run_one_to_many_encode_decode(
stage, lambda: tf.zeros([50]))
self.assertAllEqual(np.zeros((50)).astype(np.float32), test_data.decoded_x)
def test_commutes_with_sum_given_min_max(self):
stage = stages_impl.UniformQuantizationEncodingStage(bits=8,
min_max=[-1.0, 1.0])
input_values = self.evaluate([tf.random.normal([50]) for _ in range(3)])
server_test_data, decode_params = self.run_many_to_one_encode_decode(
stage, input_values)
self.assert_commutes_with_sum(
server_test_data,
stage,
decode_params,
shape=input_values[0].shape)
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_quantization_bits_deterministic_rounding(self, bits, min_max):
stage = stages_impl.UniformQuantizationEncodingStage(
bits=bits, min_max=min_max, stochastic=False)
test_data = self.run_one_to_many_encode_decode(stage, self.default_input)
self._assert_is_integer_float(test_data.encoded_x[
stages_impl.UniformQuantizationEncodingStage.ENCODED_VALUES_KEY])
# For deterministic rounding, the potential error incurred by quantization
# is bounded by half of the range of the input values divided by the number
# of quantization buckets.
if min_max is None:
self.assertAllClose(
test_data.x, test_data.decoded_x, rtol=0.0, atol=1 / (2**bits - 1))
else:
self.assertAllClose(
np.clip(test_data.x, -0.5, 0.5),
test_data.decoded_x,
rtol=0.0,
atol=0.5 / (2**bits - 1))
def test_quantization_empirically_unbiased(self):
# Tests that the quantization with stochastic=True "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 = stages_impl.UniformQuantizationEncodingStage(
bits=2, stochastic=True)
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_stochastic_tensor_raises(self):
with self.assertRaisesRegexp(TypeError, 'stochastic'):
stages_impl.UniformQuantizationEncodingStage(
stochastic=tf.constant(True, dtype=tf.bool))
def test_bad_min_max_python_values_raises(self, bad_min_max):
with self.assertRaisesRegexp(ValueError, 'smaller than the second'):
stages_impl.UniformQuantizationEncodingStage(min_max=bad_min_max)
def test_bad_min_max_python_shape_raises(self, bad_min_max):
with self.assertRaisesRegexp(ValueError, 'list with two elements'):
stages_impl.UniformQuantizationEncodingStage(min_max=bad_min_max)
def test_bad_min_max_tensor_raises(self, bad_min_max):
with self.assertRaisesRegexp(ValueError, r'shape \(2\)'):
stages_impl.UniformQuantizationEncodingStage(
min_max=tf.constant(bad_min_max))
def test_bits_out_of_range_raises(self, bits):
with self.assertRaisesRegexp(ValueError, 'integer between 1 and 16'):
stages_impl.UniformQuantizationEncodingStage(bits=bits)
def default_encoding_stage(self): """See base class.""" return stages_impl.UniformQuantizationEncodingStage()
