def testSerialization(self):
num_bits_weight = 4
num_bits_activation = 4
quantize_config = configs.DefaultNBitQuantizeConfig(
['kernel'], ['activation'], False, num_bits_weight, num_bits_activation)
expected_config = {
'class_name': 'DefaultNBitQuantizeConfig',
'config': {
'weight_attrs': ['kernel'],
'activation_attrs': ['activation'],
'quantize_output': False,
'num_bits_weight': 4,
'num_bits_activation': 4
}
}
serialized_quantize_config = tf.keras.utils.serialize_keras_object(
quantize_config)
self.assertEqual(expected_config, serialized_quantize_config)
quantize_config_from_config = tf.keras.utils.deserialize_keras_object(
serialized_quantize_config,
module_objects=globals(),
custom_objects=configs._types_dict())
self.assertEqual(quantize_config, quantize_config_from_config)
def build(self, input_shape):
conv2d_quantized = _quantize_wrapped_layer(
tf.keras.layers.Conv2D,
configs.DefaultNBitConvQuantizeConfig(
['kernel'], ['activation'],
False,
num_bits_weight=self._num_bits_weight,
num_bits_activation=self._num_bits_activation))
conv2d_quantized_output_quantized = _quantize_wrapped_layer(
tf.keras.layers.Conv2D,
configs.DefaultNBitConvQuantizeConfig(
['kernel'], ['activation'],
True,
num_bits_weight=self._num_bits_weight,
num_bits_activation=self._num_bits_activation))
num_reduced_filters = nn_layers.make_divisible(
max(1, int(self._in_filters * self._se_ratio)),
divisor=self._divisible_by)
self._se_reduce = conv2d_quantized(
filters=num_reduced_filters,
kernel_size=1,
strides=1,
padding='same',
use_bias=True,
kernel_initializer=self._kernel_initializer,
kernel_regularizer=self._kernel_regularizer,
bias_regularizer=self._bias_regularizer,
activation=NoOpActivation())
self._se_expand = conv2d_quantized_output_quantized(
filters=self._out_filters,
kernel_size=1,
strides=1,
padding='same',
use_bias=True,
kernel_initializer=self._kernel_initializer,
kernel_regularizer=self._kernel_regularizer,
bias_regularizer=self._bias_regularizer,
activation=NoOpActivation())
self._multiply = tfmot.quantization.keras.QuantizeWrapperV2(
tf.keras.layers.Multiply(),
configs.DefaultNBitQuantizeConfig(
[], [],
True,
num_bits_weight=self._num_bits_weight,
num_bits_activation=self._num_bits_activation))
self._reduce_mean_quantizer = (
tfmot.quantization.keras.quantizers.MovingAverageQuantizer(
num_bits=self._num_bits_activation,
per_axis=False,
symmetric=False,
narrow_range=False)) # activation/output
self._reduce_mean_quantizer_vars = self._reduce_mean_quantizer.build(
None, 'reduce_mean_quantizer_vars', self)
super().build(input_shape)
def testGetsResultQuantizers_EmptyWhenFalse(self):
layer = self._simple_dense_layer()
num_bits_weight = 4
num_bits_activation = 4
quantize_config = configs.DefaultNBitQuantizeConfig(
[], [], False, num_bits_weight, num_bits_activation)
output_quantizers = quantize_config.get_output_quantizers(layer)
self.assertEqual([], output_quantizers)
def testGetsResultQuantizers_ReturnsQuantizer(self):
layer = self._simple_dense_layer()
num_bits_weight = 4
num_bits_activation = 4
quantize_config = configs.DefaultNBitQuantizeConfig(
[], [], True, num_bits_weight, num_bits_activation)
output_quantizers = quantize_config.get_output_quantizers(layer)
self.assertLen(output_quantizers, 1)
self._assert_activation_quantizers(output_quantizers)
def testSetsQuantizeWeights_ErrorOnWrongShapeOfWeight(self):
layer = self._simple_dense_layer()
quantize_kernel = tf.keras.backend.variable(np.ones([1, 2]))
num_bits_weight = 4
num_bits_activation = 4
quantize_config = configs.DefaultNBitQuantizeConfig(
['kernel'], ['activation'], False, num_bits_weight, num_bits_activation)
with self.assertRaises(ValueError):
quantize_config.set_quantize_weights(layer, [quantize_kernel])
def testSetsQuantizeActivations(self):
layer = self._simple_dense_layer()
quantize_activation = tf.keras.activations.relu
num_bits_weight = 4
num_bits_activation = 4
quantize_config = configs.DefaultNBitQuantizeConfig(
['kernel'], ['activation'], False, num_bits_weight, num_bits_activation)
quantize_config.set_quantize_activations(layer, [quantize_activation])
self.assertEqual(layer.activation, quantize_activation)
def testGetsQuantizeActivationsAndQuantizers(self):
layer = self._simple_dense_layer()
num_bits_weight = 4
num_bits_activation = 4
quantize_config = configs.DefaultNBitQuantizeConfig(
['kernel'], ['activation'], False, num_bits_weight, num_bits_activation)
(activations, activation_quantizers) = self._convert_list(
quantize_config.get_activations_and_quantizers(layer))
self._assert_activation_quantizers(activation_quantizers)
self.assertEqual([layer.activation], activations)
def testSetsQuantizeWeights(self):
layer = self._simple_dense_layer()
quantize_kernel = tf.keras.backend.variable(
np.ones(layer.kernel.shape.as_list()))
num_bits_weight = 4
num_bits_activation = 4
quantize_config = configs.DefaultNBitQuantizeConfig(
['kernel'], ['activation'], False, num_bits_weight, num_bits_activation)
quantize_config.set_quantize_weights(layer, [quantize_kernel])
self._assert_kernel_equality(layer.kernel, quantize_kernel)
def testSetsQuantizeActivations_ErrorOnWrongNumberOfActivations(self):
layer = self._simple_dense_layer()
quantize_activation = tf.keras.activations.relu
num_bits_weight = 4
num_bits_activation = 4
quantize_config = configs.DefaultNBitQuantizeConfig(
['kernel'], ['activation'], False, num_bits_weight, num_bits_activation)
with self.assertRaises(ValueError):
quantize_config.set_quantize_activations(layer, [])
with self.assertRaises(ValueError):
quantize_config.set_quantize_activations(
layer, [quantize_activation, quantize_activation])
def build(self, input_shape: Optional[Union[Sequence[int], tf.Tensor]]):
"""Build variables and child layers to prepare for calling."""
conv2d_quantized = _quantize_wrapped_layer(
tf.keras.layers.Conv2D,
configs.DefaultNBitConvQuantizeConfig(
['kernel'], ['activation'], False,
num_bits_weight=self._num_bits_weight,
num_bits_activation=self._num_bits_activation))
if self._use_projection:
if self._resnetd_shortcut:
self._shortcut0 = tf.keras.layers.AveragePooling2D(
pool_size=2, strides=self._strides, padding='same')
self._shortcut1 = conv2d_quantized(
filters=self._filters * 4,
kernel_size=1,
strides=1,
use_bias=False,
kernel_initializer=self._kernel_initializer,
kernel_regularizer=self._kernel_regularizer,
bias_regularizer=self._bias_regularizer,
activation=NoOpActivation())
else:
self._shortcut = conv2d_quantized(
filters=self._filters * 4,
kernel_size=1,
strides=self._strides,
use_bias=False,
kernel_initializer=self._kernel_initializer,
kernel_regularizer=self._kernel_regularizer,
bias_regularizer=self._bias_regularizer,
activation=NoOpActivation())
self._norm0 = self._norm_with_quantize(
axis=self._bn_axis,
momentum=self._norm_momentum,
epsilon=self._norm_epsilon,
trainable=self._bn_trainable)
self._conv1 = conv2d_quantized(
filters=self._filters,
kernel_size=1,
strides=1,
use_bias=False,
kernel_initializer=self._kernel_initializer,
kernel_regularizer=self._kernel_regularizer,
bias_regularizer=self._bias_regularizer,
activation=NoOpActivation())
self._norm1 = self._norm(
axis=self._bn_axis,
momentum=self._norm_momentum,
epsilon=self._norm_epsilon,
trainable=self._bn_trainable)
self._activation1 = tfmot.quantization.keras.QuantizeWrapperV2(
tf_utils.get_activation(self._activation, use_keras_layer=True),
configs.DefaultNBitActivationQuantizeConfig(
num_bits_weight=self._num_bits_weight,
num_bits_activation=self._num_bits_activation))
self._conv2 = conv2d_quantized(
filters=self._filters,
kernel_size=3,
strides=self._strides,
dilation_rate=self._dilation_rate,
padding='same',
use_bias=False,
kernel_initializer=self._kernel_initializer,
kernel_regularizer=self._kernel_regularizer,
bias_regularizer=self._bias_regularizer,
activation=NoOpActivation())
self._norm2 = self._norm(
axis=self._bn_axis,
momentum=self._norm_momentum,
epsilon=self._norm_epsilon,
trainable=self._bn_trainable)
self._activation2 = tfmot.quantization.keras.QuantizeWrapperV2(
tf_utils.get_activation(self._activation, use_keras_layer=True),
configs.DefaultNBitActivationQuantizeConfig(
num_bits_weight=self._num_bits_weight,
num_bits_activation=self._num_bits_activation))
self._conv3 = conv2d_quantized(
filters=self._filters * 4,
kernel_size=1,
strides=1,
use_bias=False,
kernel_initializer=self._kernel_initializer,
kernel_regularizer=self._kernel_regularizer,
bias_regularizer=self._bias_regularizer,
activation=NoOpActivation())
self._norm3 = self._norm_with_quantize(
axis=self._bn_axis,
momentum=self._norm_momentum,
epsilon=self._norm_epsilon,
trainable=self._bn_trainable)
self._activation3 = tfmot.quantization.keras.QuantizeWrapperV2(
tf_utils.get_activation(self._activation, use_keras_layer=True),
configs.DefaultNBitActivationQuantizeConfig(
num_bits_weight=self._num_bits_weight,
num_bits_activation=self._num_bits_activation))
if self._se_ratio and self._se_ratio > 0 and self._se_ratio <= 1:
self._squeeze_excitation = qat_nn_layers.SqueezeExcitationNBitQuantized(
in_filters=self._filters * 4,
out_filters=self._filters * 4,
se_ratio=self._se_ratio,
kernel_initializer=self._kernel_initializer,
kernel_regularizer=self._kernel_regularizer,
bias_regularizer=self._bias_regularizer,
num_bits_weight=self._num_bits_weight,
num_bits_activation=self._num_bits_activation)
else:
self._squeeze_excitation = None
if self._stochastic_depth_drop_rate:
self._stochastic_depth = nn_layers.StochasticDepth(
self._stochastic_depth_drop_rate)
else:
self._stochastic_depth = None
self._add = tfmot.quantization.keras.QuantizeWrapperV2(
tf.keras.layers.Add(),
configs.DefaultNBitQuantizeConfig(
[], [], True,
num_bits_weight=self._num_bits_weight,
num_bits_activation=self._num_bits_activation))
super().build(input_shape)