aimet_tensorflow.keras.bn_reestimation

Top-level API

aimet_tensorflow.keras.bn_reestimation.reestimate_bn_stats(model, bn_re_estimation_dataset, bn_num_batches=100)[source]

top level api for end user directly call

Parameters:
  • model (Model) – tf.keras.Model

  • bn_re_estimation_dataset (DatasetV2) – Training dataset

  • bn_num_batches (int) – The number of batches to be used for reestimation

Return type:

Handle

Returns:

Handle that undos the effect of BN reestimation upon handle.remove()

aimet_tensorflow.keras.batch_norm_fold.fold_all_batch_norms_to_scale(sim)[source]

Fold all batch_norm layers in a model into the quantization scale parameter of the corresponding conv layers

Parameters:

sim (QuantizationSimModel) – QuantizationSimModel to be folded

Return type:

List[Tuple[QcQuantizeWrapper, QcQuantizeWrapper]]

Returns:

A list of pairs of layers [(Conv/Linear, BN layer that got folded)]