AIMET PyTorch Cross Layer Equalization APIs

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

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