def apply_quantization(model):
# Helper function uses `quantize_annotate_layer` to annotate that only the
# Dense layers should be quantized.
def add_quantize_annotation(layer):
kernelization_map = [
# tf.keras.layers.Dense,
tf.keras.layers.Conv2D
]
for layer_type in kernelization_map:
if isinstance(layer, layer_type):
quantize_config = SLCQuantizeConfig()
log.info(
"**Kernelization annotation added to layer {} of type {} with {}".format(layer.name, layer_type,
quantize_config))
quantized_layer = quantize_annotate_layer(to_annotate=layer, quantize_config=quantize_config)
return quantized_layer
log.info("**Kernelization annotation not added to layer {} of type {}".format(layer.name, type(layer)))
return layer
# Use `tf.keras.models.clone_model` to apply `add_quantize_annotation`
# to the layers of the model.
log.info("Annotating model {}".format(model.name))
annotated_model = tf.keras.models.clone_model(
model,
clone_function=add_quantize_annotation,
)
with quantize_scope({
'SLCQuantizeConfig': SLCQuantizeConfig,
"SLCWeightGenerator": SLCWeightGenerator,
"SLCRegularizer": SLCRegularizer
}):
# Use `quantize_apply` to actually make the model kernelization aware.
quant_aware_model = quantize_apply(annotated_model)
original_size = 0
compressed_size = 0
for layer in quant_aware_model.layers:
try:
original_size = original_size + layer.original_size
if layer.compressed is True:
compressed_size = compressed_size + layer.compressed_size
else:
compressed_size = compressed_size + layer.original_size
except AttributeError:
pass
try:
ratio = compressed_size * 100.0 / original_size
log.info(
"Model original size: {}, compressed size: {}, ratio: {:.2f}%".format(original_size, compressed_size, ratio))
except ZeroDivisionError:
log.info(
"Zero division error? Model original size: {}, compressed size: {}".format(original_size, compressed_size))
return quant_aware_model
def testEndToEndClusterPreserve(self):
"""Runs CQAT end to end and whole model is quantized."""
original_model = tf.keras.Sequential(
[layers.Dense(5, activation='softmax', input_shape=(10, ))])
clustered_model = cluster.cluster_weights(original_model,
**self.cluster_params)
self.compile_and_fit(clustered_model)
clustered_model = cluster.strip_clustering(clustered_model)
num_of_unique_weights_clustering = self._get_number_of_unique_weights(
clustered_model, 0, 'kernel')
quant_aware_annotate_model = (
quantize.quantize_annotate_model(clustered_model))
quant_aware_model = quantize.quantize_apply(
quant_aware_annotate_model,
scheme=default_8bit_cluster_preserve_quantize_scheme.
Default8BitClusterPreserveQuantizeScheme())
self.compile_and_fit(quant_aware_model)
stripped_cqat_model = strip_clustering_cqat(quant_aware_model)
# Check the unique weights of a certain layer of
# clustered_model and pcqat_model
num_of_unique_weights_cqat = self._get_number_of_unique_weights(
stripped_cqat_model, 1, 'kernel')
self.assertAllEqual(num_of_unique_weights_clustering,
num_of_unique_weights_cqat)
def testPassingNonPrunedModelToPCQAT(self):
"""Runs PCQAT as CQAT if the input model is not pruned."""
preserve_sparsity = False
clustered_model = self._get_clustered_model(preserve_sparsity)
clustered_model = cluster.strip_clustering(clustered_model)
nr_of_unique_weights_after = self._get_number_of_unique_weights(
clustered_model, 0, 'kernel')
# Check after plain clustering, if there are no zero weights,
# PCQAT falls back to CQAT
quant_aware_annotate_model = (
quantize.quantize_annotate_model(clustered_model))
quant_aware_model = quantize.quantize_apply(
quant_aware_annotate_model,
scheme=default_8bit_cluster_preserve_quantize_scheme.
Default8BitClusterPreserveQuantizeScheme(True))
self.compile_and_fit(quant_aware_model)
stripped_pcqat_model = strip_clustering_cqat(quant_aware_model)
# Check the unique weights of clustered_model and pcqat_model
num_of_unique_weights_pcqat = self._get_number_of_unique_weights(
stripped_pcqat_model, 1, 'kernel')
self.assertAllEqual(nr_of_unique_weights_after,
num_of_unique_weights_pcqat)
def measure_sparsity(model):
assert quantize_base.SET_CUSTOM_TNH_FLAG, log.info(
"TFMOD needs to be modified with quantizer disabled for proper "
"running")
# Helper function uses `quantize_annotate_layer` to annotate that only the
# Dense layers should be quantized.
def add_sparsity_annotation(layer):
quantize_config = SparsityMeter()
log.info(
"**Sparsity Measure annotation added to layer {} with {}".format(
layer.name, quantize_config))
quantized_layer = quantize_annotate_layer(
to_annotate=layer, quantize_config=quantize_config)
return quantized_layer
log.info("Annotating model {}".format(model.name))
tf.keras.backend.clear_session()
annotated_model = tf.keras.models.clone_model(
model, clone_function=add_sparsity_annotation)
with quantize_scope({
'SparsityMeter': SparsityMeter,
"ActivSparsityMeasure ": ActivSparsityMeasure,
"WeightsSparsityMeasure ": WeightsSparsityMeasure
}):
# Use `quantize_apply` to actually make the model Sparsity Measure aware.
quant_aware_model = quantize_apply(annotated_model)
return quant_aware_model
def cluster_preserve_quantize_model(clustered_model, train_images,
train_labels):
"""Cluster-preserve QAT model."""
quant_aware_annotate_model = (
quantize.quantize_annotate_model(clustered_model))
quant_aware_model = quantize.quantize_apply(
quant_aware_annotate_model,
scheme=default_8bit_cluster_preserve_quantize_scheme.
Default8BitClusterPreserveQuantizeScheme())
quant_aware_model.summary()
compile_and_fit(quant_aware_model, train_images, train_labels, 1)
return quant_aware_model
def apply_quantization(model):
assert quantize_base.SET_CUSTOM_TNH_FLAG, log.info(
"TFMOD needs to be modified with quantizer disabled for proper "
"running")
# Helper function uses `quantize_annotate_layer` to annotate that only the
# Dense layers should be quantized.
def add_quantize_annotation(layer):
# create new layer to break link with old model
layer = layer.__class__.from_config(layer.get_config())
quantization_map = [
# tf.keras.layers.Dense,
tf.keras.layers.Conv2D
# tf.keras.layers.Input: BFPInputQuantizerConfig()
]
for layer_type in quantization_map:
if isinstance(layer, layer_type):
quantize_config = SLGQuantizeConfig()
log.info(
"**SLG annotation added to layer {} of type {} with {}".
format(layer.name, layer_type, quantize_config))
quantized_layer = quantize_annotate_layer(
to_annotate=layer, quantize_config=quantize_config)
return quantized_layer
log.info("**SLG annotation not added to layer {} of type {}".format(
layer.name, type(layer)))
return layer
# Use `tf.keras.models.clone_model` to apply `add_quantize_annotation`
# to the layers of the model.
log.info("Annotating model {}".format(model.name))
tf.keras.backend.clear_session()
annotated_model = tf.keras.models.clone_model(
model,
clone_function=add_quantize_annotation,
)
with quantize_scope({
"SLGWeightGenerator": SLGWeightGenerator,
"SLGQuantizeConfig": SLGQuantizeConfig,
}):
# Use `quantize_apply` to actually make the model quantization aware.
quant_aware_model = quantize_apply(annotated_model)
return quant_aware_model
def testCustomWeightQuantizers_Run(self, quantizer_type):
init_params = self._get_quant_params(quantizer_type)
# Additional test that same quantizer object can be shared
# between Configs, though we don't expicitly promote this
# anywhere in the documentation.
quantizer = quantizer_type(**init_params)
class DenseQuantizeConfig(QuantizeConfig):
"""Custom QuantizeConfig for Dense layer."""
def get_weights_and_quantizers(self, layer):
return [(layer.kernel, quantizer)]
def get_activations_and_quantizers(self, layer):
# Defaults.
return [(layer.activation,
MovingAverageQuantizer(
num_bits=8,
per_axis=False,
symmetric=False,
narrow_range=False))]
def set_quantize_weights(self, layer, quantize_weights):
layer.kernel = quantize_weights[0]
def set_quantize_activations(self, layer, quantize_activations):
return
def get_output_quantizers(self, layer):
return []
def get_config(self):
return {}
annotated_model = tf.keras.Sequential([
quantize.quantize_annotate_layer(
l.Dense(8, input_shape=(10,)), DenseQuantizeConfig()),
quantize.quantize_annotate_layer(
l.Dense(5), DenseQuantizeConfig())
])
with quantize.quantize_scope(
{'DenseQuantizeConfig': DenseQuantizeConfig}):
quant_model = quantize.quantize_apply(annotated_model)
# Check no error happens.
self._train_model(quant_model)
def prune_cluster_preserve_quantize_model(clustered_model, preserve_sparsity):
"""Prune_cluster_preserve QAT model."""
pcqat_epoch = 1
quant_aware_annotate_model = quantize.quantize_annotate_model(
clustered_model)
quant_aware_model = quantize.quantize_apply(
quant_aware_annotate_model,
scheme=default_8bit_cluster_preserve_quantize_scheme.
Default8BitClusterPreserveQuantizeScheme(preserve_sparsity))
callbacks = []
quant_aware_model = _train_model(quant_aware_model, callbacks, pcqat_epoch)
pcqat_stripped = cluster_utils.strip_clustering_cqat(quant_aware_model)
return quant_aware_model, pcqat_stripped
def apply_quantization(model):
# Helper function uses `quantize_annotate_layer` to annotate that only the
# Dense layers should be quantized.
def add_quantize_annotation(layer):
# create new layer to break link with old model
layer = layer.__class__.from_config(layer.get_config())
quantization_map = {
tf.keras.layers.Dense: BFPQuantizeConfig(),
tf.keras.layers.Conv2D: BFPQuantizeConfig()
}
for layer_type, quantize_config in quantization_map.items():
if isinstance(layer, layer_type):
print(
"**Quantization annotation added to layer {} of type {} with {}"
.format(layer.name, layer_type, quantize_config))
quantized_layer = quantize_annotate_layer(
to_annotate=layer, quantize_config=quantize_config)
return quantized_layer
print("**Quantization annotation not added to layer {} of type {}".
format(layer.name, type(layer)))
return layer
# Use `tf.keras.models.clone_model` to apply `add_quantize_annotation`
# to the layers of the model.
print("Annotating model {}".format(model.name))
annotated_model = tf.keras.models.clone_model(
model,
clone_function=add_quantize_annotation,
)
with quantize_scope({
'BFPQuantizeConfig': BFPQuantizeConfig,
"BFPActivQuantizer": BFPActivQuantizer,
"BFPWeightQuantizer": BFPWeightQuantizer,
"BFPBiasQuantizer": BFPBiasQuantizer,
"PolynomialDecay": PolynomialDecay
}):
# Use `quantize_apply` to actually make the model quantization aware.
quant_aware_model = quantize_apply(annotated_model)
return quant_aware_model
def _pcqat_training(self, preserve_sparsity, quant_aware_annotate_model):
"""PCQAT training on the input model."""
quant_aware_model = quantize.quantize_apply(
quant_aware_annotate_model,
scheme=default_8bit_cluster_preserve_quantize_scheme.
Default8BitClusterPreserveQuantizeScheme(preserve_sparsity))
self.compile_and_fit(quant_aware_model)
stripped_pcqat_model = strip_clustering_cqat(quant_aware_model)
# Check the unique weights of clustered_model and pcqat_model
# layer 0 is the quantize_layer
num_of_unique_weights_pcqat = self._get_number_of_unique_weights(
stripped_pcqat_model, 1, 'kernel')
sparsity_pcqat = self._get_sparsity(stripped_pcqat_model)
return sparsity_pcqat, num_of_unique_weights_pcqat
def testEndToEndClusterPreserveOneLayer(self):
"""Runs CQAT end to end and model is quantized only for a single layer."""
original_model = tf.keras.Sequential([
layers.Dense(5, activation='relu', input_shape=(10, )),
layers.Dense(5,
activation='softmax',
input_shape=(10, ),
name='qat')
])
clustered_model = cluster.cluster_weights(original_model,
**self.cluster_params)
self.compile_and_fit(clustered_model)
clustered_model = cluster.strip_clustering(clustered_model)
num_of_unique_weights_clustering = self._get_number_of_unique_weights(
clustered_model, 1, 'kernel')
def apply_quantization_to_dense(layer):
if isinstance(layer, tf.keras.layers.Dense):
if layer.name == 'qat':
return quantize.quantize_annotate_layer(layer)
return layer
quant_aware_annotate_model = tf.keras.models.clone_model(
clustered_model,
clone_function=apply_quantization_to_dense,
)
quant_aware_model = quantize.quantize_apply(
quant_aware_annotate_model,
scheme=default_8bit_cluster_preserve_quantize_scheme.
Default8BitClusterPreserveQuantizeScheme())
self.compile_and_fit(quant_aware_model)
stripped_cqat_model = strip_clustering_cqat(quant_aware_model)
# Check the unique weights of a certain layer of
# clustered_model and pcqat_model
num_of_unique_weights_cqat = self._get_number_of_unique_weights(
stripped_cqat_model, 1, 'kernel')
self.assertAllEqual(num_of_unique_weights_clustering,
num_of_unique_weights_cqat)
def testModelEndToEnd(self, model_type):
# 1. Check whether quantized model graph can be constructed.
model = self._get_model(model_type)
model = quantize.quantize_apply(quantize.quantize_annotate(model))
# 2. Sanity check to ensure basic training on random data works.
x_train, y_train = self._create_test_data(model)
model.compile(loss='categorical_crossentropy',
optimizer='sgd',
metrics=['accuracy'])
model.fit(x_train, y_train)
# 3. Ensure conversion to TFLite works.
_, tflite_file = tempfile.mkstemp('.tflite')
print('TFLite File: ', tflite_file)
with quantize.quantize_scope():
utils.convert_keras_to_tflite(
model,
tflite_file,
inference_input_type=tf.lite.constants.FLOAT)
# 4. Verify input runs on converted model.
self._verify_tflite(tflite_file, x_train, y_train)
def prune_preserve_quantize_model(pruned_model, train_images, train_labels):
batch_size = 256
epochs = 5
pruned_model = prune.strip_pruning(pruned_model)
# Prune preserve QAT model
quant_aware_annotate_model = quantize.quantize_annotate_model(pruned_model)
quant_aware_model = quantize.quantize_apply(
quant_aware_annotate_model,
scheme=default_8bit_prune_preserve_quantize_scheme
.Default8BitPrunePreserveQuantizeScheme())
quant_aware_model.summary()
fit_kwargs = {
'batch_size': batch_size,
'epochs': epochs,
}
compile_and_fit(quant_aware_model,
train_images,
train_labels,
compile_kwargs={},
fit_kwargs=fit_kwargs)
return quant_aware_model
padding='same',
activation='relu'),
input_shape=input_shape),
l.MaxPooling2D((2, 2), (2, 2), padding='same'),
quantize.quantize_annotate_layer(
l.Conv2D(64, 5, padding='same', activation='relu')),
l.MaxPooling2D((2, 2), (2, 2), padding='same'),
l.Flatten(),
quantize.quantize_annotate_layer(l.Dense(1024, activation='relu')),
l.Dropout(0.4),
quantize.quantize_annotate_layer(l.Dense(num_classes)),
# TODO(alanchiao): fuse softmax once we've handled it.
l.Softmax(),
])
model = quantize.quantize_apply(model)
# Dump graph to /tmp for verification on tensorboard.
graph_def = tf.get_default_graph().as_graph_def()
with open('/tmp/mnist_model.pbtxt', 'w') as f:
f.write(str(graph_def))
model.compile(loss=tf.keras.losses.categorical_crossentropy,
optimizer=tf.keras.optimizers.Adadelta(),
metrics=['accuracy'])
model.fit(x_train,
y_train,
batch_size=batch_size,
epochs=epochs,
verbose=1,
def testTrainableWeightsBehaveCorrectlyDuringPCQAT(self):
"""PCQAT zero centroid masks stay the same and trainable variables are updating between epochs."""
preserve_sparsity = True
clustered_model = self._get_clustered_model(preserve_sparsity)
clustered_model = cluster.strip_clustering(clustered_model)
# Apply PCQAT
quant_aware_annotate_model = (
quantize.quantize_annotate_model(clustered_model))
quant_aware_model = quantize.quantize_apply(
quant_aware_annotate_model,
scheme=default_8bit_cluster_preserve_quantize_scheme.
Default8BitClusterPreserveQuantizeScheme(True))
quant_aware_model.compile(
loss=tf.keras.losses.categorical_crossentropy,
optimizer='adam',
metrics=['accuracy'],
)
class CheckCentroidsAndTrainableVarsCallback(
tf.keras.callbacks.Callback):
"""Check the updates of trainable variables and centroid masks."""
def on_epoch_begin(self, batch, logs=None):
# Check cluster centroids have the zero in the right position
vars_dictionary = self.model.layers[1]._weight_vars[0][2]
self.centroid_mask = vars_dictionary['centroids_mask']
self.zero_centroid_index_begin = np.where(
self.centroid_mask == 0)[0]
# Check trainable weights before training
self.layer_kernel = (self.model.layers[1].weights[3].numpy())
self.original_weight = vars_dictionary[
'ori_weights_vars_tf'].numpy()
self.centroids = vars_dictionary['cluster_centroids_tf'].numpy(
)
def on_epoch_end(self, batch, logs=None):
# Check the index of the zero centroids are not changed after training
vars_dictionary = self.model.layers[1]._weight_vars[0][2]
self.zero_centroid_index_end = np.where(
vars_dictionary['centroids_mask'] == 0)[0]
assert np.array_equal(self.zero_centroid_index_begin,
self.zero_centroid_index_end)
# Check trainable variables after training are updated
assert not np.array_equal(
self.layer_kernel, self.model.layers[1].weights[3].numpy())
assert not np.array_equal(
self.original_weight,
vars_dictionary['ori_weights_vars_tf'].numpy())
assert not np.array_equal(
self.centroids,
vars_dictionary['cluster_centroids_tf'].numpy())
# Use many epochs to verify layer's kernel weights are updating because
# they can stay the same after being trained using only the first batch
# of data for instance
quant_aware_model.fit(
np.random.rand(20, 10),
tf.keras.utils.to_categorical(np.random.randint(5, size=(20, 1)),
5),
steps_per_epoch=5,
epochs=3,
callbacks=[CheckCentroidsAndTrainableVarsCallback()])
def apply_quantization(model,
pruning_policy=None,
weight_precision=None,
activation_precision=None,
activation_margin=None):
# assert quantize_base.SET_CUSTOM_TNH_FLAG, log.info("TFMOD needs to be modified with quantizer disabled for proper "
# "running")
if weight_precision is not None:
global _WEIGHTS_NUM_BITS # need to declare when you want to change the value
_WEIGHTS_NUM_BITS = weight_precision
if activation_precision is not None:
global _ACTIV_NUM_BITS
_ACTIV_NUM_BITS = activation_precision
if activation_margin is not None:
global _ACTIV_MARGIN
_ACTIV_MARGIN = activation_margin
log.info(
"Weights num bits: {} - Activ num bits: {} - Activ margin: {}".format(
_WEIGHTS_NUM_BITS, _ACTIV_NUM_BITS, _ACTIV_MARGIN))
# Helper function uses `quantize_annotate_layer` to annotate that only the
# Dense layers should be quantized.
def add_quantize_annotation(layer):
# create new layer to break link with old model
try:
layer = layer.__class__.from_config(layer.get_config())
except:
pass
for layer_type in quantization_map:
if isinstance(layer, layer_type):
if isinstance(pruning_policy, float) or pruning_policy is None:
layer_pruning = pruning_policy
elif isinstance(pruning_policy, dict):
layer_pruning = pruning_policy[layer.name]
else:
raise ValueError("Illegal layer pruning policy {}".format(
pruning_policy))
quantize_config = BFPQuantizeConfig(
pruning_policy=layer_pruning)
log.info(
"**Quantization annotation added to layer {} of type {} with {}"
.format(layer.name, layer_type, quantize_config))
quantized_layer = quantize_annotate_layer(
to_annotate=layer, quantize_config=quantize_config)
return quantized_layer
log.info("**Quantization annotation not added to layer {} of type {}".
format(layer.name, type(layer)))
return layer
# Use `tf.keras.models.clone_model` to apply `add_quantize_annotation`
# to the layers of the model.
log.info("Annotating model {}".format(model.name))
tf.keras.backend.clear_session()
annotated_model = tf.keras.models.clone_model(
model,
clone_function=add_quantize_annotation,
)
with quantize_scope({
'BFPQuantizeConfig': BFPQuantizeConfig,
"BFPActivQuantizer": BFPActivQuantizer,
"BFPWeightQuantizer": BFPWeightQuantizer,
"BFPBiasQuantizer": BFPBiasQuantizer,
"PolynomialDecay": PolynomialDecay
}):
# Use `quantize_apply` to actually make the model quantization aware.
quant_aware_model = quantize_apply(annotated_model)
for q_layer in quant_aware_model.layers:
if isinstance(q_layer, QuantizeWrapper):
for quant_type in quantization_map:
if isinstance(q_layer.layer, quant_type):
original_name = q_layer.name.replace("quant_", "")
old_layer = model.get_layer(original_name)
q_weights = q_layer.get_weights()
orig_weights = old_layer.get_weights()
q_weights[0] = orig_weights[0]
try:
q_weights[1] = orig_weights[1]
except IndexError:
pass
q_layer.set_weights(q_weights)
return quant_aware_model
def quantize(self):
model = self.get_model()
def apply_quantization_to_dense(layer):
if isinstance(layer, tf.keras.layers.Dense) or isinstance(
layer, tf.keras.layers.Conv2D):
return quantize_annotate_layer(layer)
return layer
# Use `tf.keras.models.clone_model` to apply `apply_quantization_to_dense`
# to the layers of the model.
annotated_model = tf.keras.models.clone_model(
model,
clone_function=apply_quantization_to_dense,
)
# Now that the Dense layers are annotated,
# `quantize_apply` actually makes the model quantization aware.
quant_aware_model = quantize_apply(annotated_model)
quant_aware_model.summary()
(train_images,
train_labels), (test_images,
test_labels) = self.get_dataset().load_data()
self.compile_model(quant_aware_model)
self.compile_model(model)
CLASS_NAMES = [
'airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog',
'horse', 'ship', 'truck'
]
for x, y in tf.data.Dataset.from_tensor_slices(
(test_images, test_labels)).map(self.process_images).shuffle(
buffer_size=1024).take(10).batch(1):
plt.title(CLASS_NAMES[y[0][0]])
plt.imshow(x[0])
print([
CLASS_NAMES[i]
for i in tf.argsort(quant_aware_model.predict(x))[0]
])
print([CLASS_NAMES[i] for i in tf.argsort(model.predict(x))[0]])
plt.show()
exit()
converter = tf.lite.TFLiteConverter.from_keras_model(model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
(train_images,
train_labels), (test_images,
test_labels) = self.get_dataset().load_data()
train_images, train_labels = train_images[5000:], train_labels[5000:]
def representative_dataset():
for data, label in tf.data.Dataset.from_tensor_slices(
(train_images,
train_labels)).map(self.process_images).batch(1).take(100):
yield [tf.cast(data, tf.float32)]
converter.representative_dataset = representative_dataset
converter.target_spec.supported_ops = [
tf.lite.OpsSet.TFLITE_BUILTINS_INT8
]
converter.inference_input_type = tf.int8 # or tf.uint8
converter.inference_output_type = tf.int8 # or tf.uint8
tflite_quant_model = converter.convert()
logger.info('Model quantized for tensorflow lite successfully')