Function main

examples/qualcomm/oss_scripts/fastvit.py:59–163

(args)

def main(args):
    qnn_config = QnnConfig.load_config(args.config_file if args.config_file else args)

    # ensure the working directory exist.
    os.makedirs(args.artifact, exist_ok=True)

    data_num = 100
    inputs, targets = get_imagenet_dataset(
        dataset_path=f"{args.dataset}",
        data_size=data_num,
        image_shape=(256, 256),
    )

    pte_filename = "fastvit_qnn"

    def get_custom_quantizer():
        quantizer = make_quantizer(
            quant_dtype=QuantDtype.use_8a8w,
            backend=qnn_config.backend,
            soc_model=qnn_config.soc_model,
        )

        # there are lots of outliers appearing in fastvit parameters
        # we need to apply special configuration to saturate their impact
        act_qspec = QuantizationSpec(
            dtype=torch.uint8,
            qscheme=torch.per_tensor_affine,
            observer_or_fake_quant_ctr=MovingAverageMinMaxObserver.with_args(
                **{"averaging_constant": 0.01}
            ),
        )
        weight_qspec = QuantizationSpec(
            dtype=torch.int8,
            quant_min=torch.iinfo(torch.int8).min + 1,
            quant_max=torch.iinfo(torch.int8).max,
            qscheme=torch.per_channel_symmetric,
            ch_axis=0,
            observer_or_fake_quant_ctr=PerChannelParamObserverWithLossEvaluation.with_args(
                **{"steps": 100, "use_mse": True}
            ),
        )
        # rewrite default per-channel ptq config
        quantizer.default_quant_config.per_channel_quant_config = QuantizationConfig(
            input_activation=act_qspec,
            output_activation=act_qspec,
            weight=weight_qspec,
            bias=_derived_bias_quant_spec,
        )

        # rewrite default ptq config
        q_config = quantizer.default_quant_config.quant_config
        quantizer.default_quant_config.quant_config = QuantizationConfig(
            input_activation=act_qspec,
            output_activation=act_qspec,
            weight=q_config.weight,
            bias=q_config.bias,
        )
        return quantizer