AIMET PyTorch Cross Layer Equalization APIs

User Guide Link

To learn more about this technique, please see Cross-Layer Equalization

Examples Notebook Link

For an end-to-end notebook showing how to use PyTorch Cross Layer Equalization, please see here.

Introduction

AIMET functionality for PyTorch Cross Layer Equalization has 3 features-

• BatchNorm Folding

• Cross Layer Scaling

• High Bias Fold

Cross Layer Equalization API

The following API performs BatchNorm fold followed by Cross Layer Scaling followed by High Bias Fold.

Note: High Bias fold will not happen when the below API is used, if the model does not have BatchNorm layers

API for Cross Layer Equalization

aimet_torch.cross_layer_equalization.equalize_model(model, input_shapes, dummy_input=None)

High-level API to perform Cross-Layer Equalization (CLE) on the given model. The model is equalized in place.

Parameters

• model (Module) – Model to equalize

• input_shapes (Union[Tuple, List[Tuple]]) – Shape of the input (can be a tuple or a list of tuples if multiple inputs)

• dummy_input (Union[Tensor, Tuple, None]) – A dummy input to the model. Can be a Tensor or a Tuple of Tensors

Returns

None

Code Example

Required imports

from torchvision import models
from aimet_torch.cross_layer_equalization import equalize_model

Cross Layer Equalization in auto mode

def cross_layer_equalization_auto():
    model = models.resnet18(pretrained=True)

    input_shape = (1, 3, 224, 224)

    model = model.eval()

    # Performs BatchNorm fold, Cross layer scaling and High bias folding
    equalize_model(model, input_shape)

Primitive APIs

If the user would like to call the APIs individually, then the following APIs can be used-

• Primitive APIs for Cross Layer Equalization