def _runQuantizeTest(self, range_min, range_max, data_dtype,
quantization_dtype, expected_scale):
d = np.arange(range_min, range_max + 1, dtype=data_dtype)
q, s, m = quantization.quantize_weights(d, quantization_dtype)
self.assertAlmostEqual(s, expected_scale)
self.assertEqual(q.dtype, quantization_dtype)
de_q = quantization.dequantize_weights(q, s, m, data_dtype)
np.testing.assert_allclose(de_q, d)
if range_min <= 0 <= range_max:
d_0 = np.zeros(1, data_dtype)
q_0 = np.round((d_0 - m) / s).astype(quantization_dtype)
self.assertEqual(
quantization.dequantize_weights(q_0, s, m, data_dtype), d_0)
def _runQuantizeTest(
self, range_min, range_max, data_dtype, quantization_dtype,
expected_scale):
d = np.arange(range_min, range_max + 1, dtype=data_dtype)
q, s, m = quantization.quantize_weights(d, quantization_dtype)
self.assertAlmostEqual(s, expected_scale)
self.assertEqual(q.dtype, quantization_dtype)
de_q = quantization.dequantize_weights(q, s, m, data_dtype)
np.testing.assert_allclose(de_q, d)
if range_min <= 0 <= range_max:
d_0 = np.zeros(1, data_dtype)
q_0 = np.round((d_0 - m) / s).astype(quantization_dtype)
self.assertEqual(
quantization.dequantize_weights(q_0, s, m, data_dtype), d_0)
def testAllEqual(self):
d = np.ones(5, dtype=np.float32)
q, s, m = quantization.quantize_weights(d, np.uint8)
self.assertEqual(s, 1.0)
self.assertEqual(q.dtype, np.uint8)
de_q = quantization.dequantize_weights(q, s, m, np.float32)
np.testing.assert_array_equal(de_q, d)
def testAllEqual(self):
d = np.ones(5, dtype=np.float32)
q, s, m = quantization.quantize_weights(d, np.uint8)
self.assertEqual(s, 1.0)
self.assertEqual(q.dtype, np.uint8)
de_q = quantization.dequantize_weights(q, s, m, np.float32)
np.testing.assert_array_equal(de_q, d)
def testFloatQuantizeAllEqual(self):
d = np.ones(5, dtype=np.float32)
q, metadata = quantization.quantize_weights(d, np.float16)
self.assertDictEqual(metadata, {})
self.assertEqual(q.dtype, np.float16)
de_q = quantization.dequantize_weights(q, metadata, np.float32)
np.testing.assert_array_equal(de_q, d)
def testAffineQuantizeAllEqual(self):
d = np.ones(5, dtype=np.float32)
q, metadata = quantization.quantize_weights(d, np.uint8)
assert 'scale' in metadata and 'min' in metadata
self.assertEqual(metadata['scale'], 1.0)
self.assertEqual(q.dtype, np.uint8)
de_q = quantization.dequantize_weights(q, metadata, np.float32)
np.testing.assert_array_equal(de_q, d)
def _runQuantizeTest(self, range_min, range_max, data_dtype,
quantization_dtype, expected_metadata):
d = np.arange(range_min, range_max + 1, dtype=data_dtype)
q, metadata = quantization.quantize_weights(d, quantization_dtype)
self.assertDictContainsSubsetAlmostEqual(metadata, expected_metadata)
self.assertEqual(q.dtype, quantization_dtype)
de_q = quantization.dequantize_weights(q, metadata, data_dtype)
np.testing.assert_allclose(de_q, d)
if quantization_dtype in [np.uint8, np.uint16]:
s = metadata['scale']
m = metadata['min']
if range_min <= 0 <= range_max:
d_0 = np.zeros(1, data_dtype)
q_0 = np.round((d_0 - m) / s).astype(quantization_dtype)
self.assertEqual(
quantization.dequantize_weights(q_0, metadata, data_dtype),
d_0)
def testInvalidDequantizationTypes(self):
# Invalid metadata for affine quantization
with self.assertRaises(ValueError):
d = np.ones(1, dtype=np.uint8)
quantization.dequantize_weights(np.array([]), {})
# Invalid target dtype for float16 quantization
with self.assertRaises(ValueError):
d = np.ones(1, dtype=np.float16)
quantization.dequantize_weights(d, {}, np.int32)
# Invalid dequantization type
with self.assertRaises(ValueError):
d = np.ones(1, dtype=np.bool)
quantization.dequantize_weights(d, {})
def decode_weights(weights_manifest, data_buffers, flatten=False):
"""Load weight values from buffer(s) according to a weights manifest.
Args:
weights_manifest: A TensorFlow.js-format weights manifest (a JSON array).
data_buffers: A buffer or a `list` of buffers containing the weights values
in binary format, concatenated in the order specified in
`weights_manifest`. If a `list` of buffers, the length of the `list`
must match the length of `weights_manifest`. A single buffer is
interpreted as a `list` of one buffer and is valid only if the length of
`weights_manifest` is `1`.
flatten: Whether all the weight groups in the return value are to be
flattened as a single weight groups. Default: `False`.
Returns:
If `flatten` is `False`, a `list` of weight groups. Each group is an array
of weight entries. Each entry is a dict that maps a unique name to a numpy
array, for example:
entry = {
'name': 'weight1',
'data': np.array([1, 2, 3], 'float32')
}
Weights groups would then look like:
weight_groups = [
[group_0_entry1, group_0_entry2],
[group_1_entry1, group_1_entry2],
]
If `flatten` is `True`, returns a single weight group.
Raises:
ValueError: if the lengths of `weights_manifest` and `data_buffers` do not
match.
"""
if not isinstance(data_buffers, list):
data_buffers = [data_buffers]
if len(weights_manifest) != len(data_buffers):
raise ValueError(
'Mismatch in the length of weights_manifest (%d) and the length of '
'data buffers (%d)' % (len(weights_manifest), len(data_buffers)))
out = []
for group, data_buffer in zip(weights_manifest, data_buffers):
offset = 0
out_group = []
for weight in group['weights']:
quant_info = weight.get('quantization', None)
name = weight['name']
if weight['dtype'] == 'string':
# String array.
dtype = np.object
elif quant_info:
# Quantized array.
dtype = np.dtype(quant_info['dtype'])
else:
# Regular numeric array.
dtype = np.dtype(weight['dtype'])
shape = weight['shape']
if dtype not in _INPUT_DTYPES:
raise NotImplementedError('Unsupported data type: %s' % dtype)
if weight['dtype'] == 'string':
value, offset = _deserialize_string_array(data_buffer, offset, shape)
else:
value = _deserialize_numeric_array(data_buffer, offset, dtype, shape)
offset += dtype.itemsize * value.size
if quant_info:
value = quantization.dequantize_weights(
value, quant_info, np.dtype(weight['dtype']))
out_group.append({'name': name, 'data': value})
if flatten:
out += out_group
else:
out.append(out_group)
return out
def testAffineQuantizeNormalizedFloats(self):
data = np.array([-0.29098126, -0.24776903, -0.27248842, 0.23848203],
dtype=np.float32)
q, metadata = quantization.quantize_weights(data, np.uint16)
de_q = quantization.dequantize_weights(q, metadata, data.dtype)
np.testing.assert_array_almost_equal(de_q, data, decimal=5)
def decode_weights(weights_manifest, data_buffers, flatten=False):
"""Load weight values from buffer(s) according to a weights manifest.
Args:
weights_manifest: A TensorFlow.js-format weights manifest (a JSON array).
data_buffers: A buffer or a `list` of buffers containing the weights values
in binary format, concatenated in the order specified in
`weights_manifest`. If a `list` of buffers, the length of the `list`
must match the length of `weights_manifest`. A single buffer is
interpreted as a `list` of one buffer and is valid only if the length of
`weights_manifest` is `1`.
flatten: Whether all the weight groups in the return value are to be
flattened as a single weight groups. Default: `False`.
Returns:
If `flatten` is `False`, a `list` of weight groups. Each group is an array
of weight entries. Each entry is a dict that maps a unique name to a numpy
array, for example:
entry = {
'name': 'weight1',
'data': np.array([1, 2, 3], 'float32')
}
Weights groups would then look like:
weight_groups = [
[group_0_entry1, group_0_entry2],
[group_1_entry1, group_1_entry2],
]
If `flatten` is `True`, returns a single weight group.
Raises:
ValueError: if the lengths of `weights_manifest` and `data_buffers` do not
match.
"""
if not isinstance(data_buffers, list):
data_buffers = [data_buffers]
if len(weights_manifest) != len(data_buffers):
raise ValueError(
'Mismatch in the length of weights_manifest (%d) and the length of '
'data buffers (%d)' % (len(weights_manifest), len(data_buffers)))
out = []
for group, data_buffer in zip(weights_manifest, data_buffers):
offset = 0
out_group = []
for weight in group['weights']:
quantization_info = weight.get('quantization', None)
name = weight['name']
dtype = np.dtype(
quantization_info['dtype'] if quantization_info else weight['dtype'])
shape = weight['shape']
if dtype not in _INPUT_DTYPES:
raise NotImplementedError('Unsupported data type: %s' % dtype)
unit_bytes = dtype.itemsize
weight_numel = 1
for dim in shape:
weight_numel *= dim
weight_bytes = unit_bytes * weight_numel
value = np.frombuffer(
data_buffer, dtype=dtype, count=weight_numel,
offset=offset).reshape(shape)
if quantization_info:
value = quantization.dequantize_weights(
value, quantization_info['scale'], quantization_info['min'],
np.dtype(weight['dtype']))
offset += weight_bytes
out_group.append({'name': name, 'data': value})
if flatten:
out += out_group
else:
out.append(out_group)
return out
