AIMET TensorFlow AutoQuant API¶
Top-level API¶
-
class
aimet_tensorflow.auto_quant.AutoQuant(allowed_accuracy_drop, unlabeled_dataset, eval_callback, default_param_bw=8, default_output_bw=8, default_quant_scheme=<QuantScheme.post_training_tf_enhanced: 2>, default_rounding_mode='nearest', default_config_file=None)¶ Integrate and apply post-training quantization techniques.
AutoQuant includes 1) batchnorm folding, 2) cross-layer equalization, and 3) Adaround. These techniques will be applied in a best-effort manner until the model meets the evaluation goal given as allowed_accuracy_drop.
- Parameters
allowed_accuracy_drop (
float) – Maximum allowed accuracy drop.unlabeled_dataset (
DatasetV1) – An unlabeled dataset for encoding computation. By default, this dataset will be also used for Adaround unless otherwise specified by self.set_adaround_params.eval_callback (
Callable[[Session,Optional[int]],float]) – A function that maps a tf session and the number of samples to the evaluation score. This callback is expected to return a scalar value representing the model performance evaluated against exactly N samples, where N is the number of samples passed as the second argument of this callback. NOTE: If N is None, the model is expected to be evaluated against the whole evaluation dataset.default_param_bw (
int) – Default bitwidth (4-31) to use for quantizing layer parameters.default_output_bw (
int) – Default bitwidth (4-31) to use for quantizing layer inputs andoutputs.default_quant_scheme (
QuantScheme) – Quantization scheme. Supported values are QuantScheme.post_training_tf or QuantScheme.post_training_tf_enhanced.default_rounding_mode (
str) – Rounding mode. Supported options are ‘nearest’ or ‘stochastic’default_config_file (
Optional[str]) – Path to configuration file for model quantizers
-
apply(fp32_sess, starting_op_names, output_op_names, results_dir='/tmp', cache_id=None)¶ Apply post-training quantization techniques.
- Parameters
fp32_sess (
Session) – tf.Session associated with the model to apply PTQ techniques.starting_op_names (
List[str]) – List of starting op names of the model.output_op_names (
List[str]) – List of output op names of the model.results_dir (
str) – Directory to save the results.
- Return type
Tuple[Session,float,str]- Returns
Tuple of (best session, eval score, encoding path).
-
set_adaround_params(adaround_params)¶ Set Adaround parameters. If this method is not called explicitly by the user, AutoQuant will use unlabeled_dataset (passed to __init__) for Adaround.
- Parameters
adaround_params (
AdaroundParameters) – Adaround parameters.- Return type
None
Code Examples¶
Required imports
from typing import Optional
import numpy as np
import tensorflow as tf
from tensorflow.keras.applications.resnet import ResNet50
from aimet_tensorflow.utils.common import iterate_tf_dataset
from aimet_tensorflow.adaround.adaround_weight import AdaroundParameters
from aimet_tensorflow.auto_quant import AutoQuant
from aimet_tensorflow.utils.graph import update_keras_bn_ops_trainable_flag
tf.compat.v1.disable_eager_execution()
Define constants and helper functions
EVAL_DATASET_SIZE = 5000
CALIBRATION_DATASET_SIZE = 2000
BATCH_SIZE = 100
_sampled_datasets = {}
def _create_sampled_dataset(dataset, num_samples):
if num_samples in _sampled_datasets:
return _sampled_datasets[num_samples]
with dataset._graph.as_default():
SHUFFLE_BUFFER_SIZE = 300 # NOTE: Adjust the buffer size as necessary.
SHUFFLE_SEED = 22222
dataset = dataset.shuffle(buffer_size=SHUFFLE_BUFFER_SIZE, seed=SHUFFLE_SEED)\
.take(num_samples)\
.batch(BATCH_SIZE)
_sampled_datasets[num_samples] = dataset
return dataset
Prepare model and dataset
input_shape = (224, 224, 3)
num_classes = 1000
model = ResNet50(weights='imagenet', input_shape=input_shape)
model = update_keras_bn_ops_trainable_flag(model, False, load_save_path='./')
input_tensor_name = model.input.name
input_op_name, _ = input_tensor_name.split(":")
output_tensor_name = model.output.name
output_op_name, _ = output_tensor_name.split(":")
# NOTE: In the actual use cases, a real dataset should provide by the users.
images = np.random.rand(100, *input_shape)
labels = np.random.randint(num_classes, size=(100,))
image_dataset = tf.compat.v1.data.Dataset.from_tensor_slices(images)\
.repeat()\
.take(EVAL_DATASET_SIZE)
label_dataset = tf.compat.v1.data.Dataset.from_tensor_slices(labels)\
.repeat()\
.take(EVAL_DATASET_SIZE)
eval_dataset = tf.compat.v1.data.Dataset.zip((image_dataset, label_dataset))
Prepare unlabeled dataset
# NOTE: In the actual use cases, the users should implement this part to serve
# their own goals if necessary.
unlabeled_dataset = image_dataset.batch(BATCH_SIZE)
Prepare eval callback
# NOTE: In the actual use cases, the users should implement this part to serve
# their own goals if necessary.
def eval_callback(sess: tf.compat.v1.Session,
num_samples: Optional[int] = None) -> float:
if num_samples is None:
num_samples = EVAL_DATASET_SIZE
sampled_dataset = _create_sampled_dataset(eval_dataset, num_samples)
with sess.graph.as_default():
sess.run(tf.compat.v1.global_variables_initializer())
input_tensor = sess.graph.get_tensor_by_name(input_tensor_name)
output_tensor = sess.graph.get_tensor_by_name(output_tensor_name)
num_correct_predictions = 0
for images, labels in iterate_tf_dataset(sampled_dataset):
prob = sess.run(output_tensor, feed_dict={input_tensor: images})
predictions = np.argmax(prob, axis=1)
num_correct_predictions += np.sum(predictions == labels)
return int(num_correct_predictions) / num_samples
Create AutoQuant object
auto_quant = AutoQuant(allowed_accuracy_drop=0.01,
unlabeled_dataset=unlabeled_dataset,
eval_callback=eval_callback)
(Optional) Set Adaround parameters
For setting the num_batches parameter, use the following guideline. The number of batches is used to evaluate the model while calculating the quantization encodings. Typically we want AdaRound to use around 2000 samples. For example, if the batch size is 32, num_batches is 64. If the batch size you are using is different, adjust the num_batches accordingly.
ADAROUND_DATASET_SIZE = 2000
adaround_dataset = _create_sampled_dataset(image_dataset, ADAROUND_DATASET_SIZE)
adaround_params = AdaroundParameters(adaround_dataset,
num_batches=ADAROUND_DATASET_SIZE // BATCH_SIZE)
auto_quant.set_adaround_params(adaround_params)
Run AutoQuant
sess, accuracy, encoding_path =\
auto_quant.apply(tf.compat.v1.keras.backend.get_session(),
starting_op_names=[input_op_name],
output_op_names=[output_op_name])