Package: tensorflow-batchnorm-folding
Batch-Normalization FoldingIn this repository, we propose an implementation of the batch-normalization folding algorithm fromIJCAI 2022. Batch-Normalization Folding implements the batch normalization layer by folding it into a appropriate layer. The original batch-normalization layers are removed without changing the predictive function defined by the neural network. The simpliest scenario is an application for a fully-connected layer followed by a batch-normalization layer, we getx \mapsto \gamma \frac{Ax + b - \mu}{\sigma + \epsilon} + \beta = \gamma \frac{A}{\sigma +\epsilon} x + \frac{b - \mu}{\sigma + \epsilon} + \betaThus the two layers can be expressed as a single fully-connected layer at inference without any change in the predictive function.UseThis repository is available as a pip package (usepip install tensorflow-batchnorm-folding).
This implementation is compatible with tf.keras.Model instances. It was tested with the following modelsResNet 50MobileNet V2MobileNet V3EfficentNet B0To run a simple test:frombatch_normalization_folding.folderimportfold_batchnormalization_layersimporttensorflowastfmod=tf.keras.applications.efficientnet.EfficientNetB0()folded_model,output_str=fold_batchnormalization_layers(mod,verbose=True)Theoutput_stris either the ratio num_layers_folded/num_layers_not_folded or 'failed' to state a failure in the process.ParametersThe functionfold_batchnormalization_layerscan be called with multiple parameters.ParameterOptionmodelModel to fold.folding_mechanismban-offuses BaN-OFF als folding mechanism (recommended)simpleuses a folding mechanism that only folds neighboring layers.Default:ban-offverboseTrueprints additional information during the folding process.FalseDisables the printing of additional information.Default:FalseLayer foldingThis Python implementation supports folding into many types of layers, supported is the folding of batch-normalization parameters into:DenseConv1DConv2DDepthwiseConv2DTo Dounit test on all keras applciations modelscheck package installementdeal with Concatenate layersCite@inproceedings{yvinec2022fold,
  title={To Fold or Not to Fold: a Necessary and Sufficient Condition on Batch-Normalization Layers Folding},
  author={Yvinec, Edouard and Dapogny, Arnaud and Bailly, Kevin},
  journal={IJCAI},
  year={2022}
}Performance on Base Models+------------------------------------+
|             ResNet 50              |
+------------------------------------+
| BN layers folded         |    53   |
| BN layers not folded     |    0    |
+------------------------------------+
|          EfficientNet B0           |
+------------------------------------+
| BN layers folded         |    49   |
| BN layers not folded     |    0    |
+------------------------------------+
|            MobileNet V2            |
+------------------------------------+
| BN layers folded         |    52   |
| BN layers not folded     |    0    |
+------------------------------------+
|            MobileNet V3            |
+------------------------------------+
| BN layers folded         |    34   |
| BN layers not folded     |    0    |
+------------------------------------+
|        Inception ResNet V2         |
+------------------------------------+
| BN layers folded         |   204   |
| BN layers not folded     |    0    |
+------------------------------------+
|            Inception V3            |
+------------------------------------+
| BN layers folded         |    94   |
| BN layers not folded     |    0    |
+------------------------------------+
|               NASNet               |
+------------------------------------+
| BN layers folded         |    28   |
| BN layers not folded     |   164   |
+------------------------------------+
|            DenseNet 121            |
+------------------------------------+
| BN layers folded         |    59   |
| BN layers not folded     |    62   |
+------------------------------------+