aimet_tensorflow.batch_norm_fold

Top-level API

aimet_tensorflow.keras.batch_norm_fold.fold_all_batch_norms(model)

Fold all batch_norm layers in a model into corresponding conv/linear layers

Parameters:

model (Model) – model to find all batch norms for

Return type:

Tuple[List[Tuple[Union[Conv2D, Dense, Conv2DTranspose, DepthwiseConv2D], BatchNormalization]], Model]

Returns:

A tuple of List of conv/linear layers with associated bn op / activation info and a new model with the Batch Normalization layers folded