Function quantize

examples/qualcomm/util_scripts/cli.py:159–203

(args)

def quantize(args):
    logger = get_logger()
    # get corresponding QnnQuantizer
    try:
        quant_dtype = getattr(QuantDtype, args.config)
        act_observer = getattr(pt2e, args.activation_observer)
        quantizer = make_quantizer(
            quant_dtype=quant_dtype,
            per_channel_conv=args.per_channel,
            per_channel_linear=args.per_row,
            act_observer=act_observer,
            backend=get_backend_type(args.backend),
            soc_model=args.soc_model,
            eps=args.eps,
        )
    except Exception:
        logger.error(
            f"Failed to retrieve expected config {args.config} / {args.activation_observer}."
        )
        exit(1)

    # step 0: load saved model
    ep = torch.export.load(args.artifact)
    # step 1: use prepare_pt2e to annotate QDQ pairs
    ep_prepared = prepare_pt2e(ep.module(), quantizer)
    logger.info(f"perform calibration on {args.artifact}")
    # step 2: perform calibration
    input_list_parser = InputListParser(args.input_list)
    graph_input_names = [
        spec.arg.name
        for spec in ep.graph_signature.input_specs
        if spec.kind.name == "USER_INPUT"
    ]
    for inputs in input_list_parser:
        if isinstance(inputs, dict):
            inputs = [inputs[name] for name in graph_input_names]
        ep_prepared(*inputs)
    # step 3: use convert_pt2e to fix encodings of QDQ pairs
    logger.info(f"saving calibrated model for {args.artifact}")
    ep_converted = convert_pt2e(ep_prepared)
    ep_quantized = torch.export.export(ep_converted, tuple(inputs))
    os.makedirs(args.output_folder, exist_ok=True)
    torch.export.save(
        ep_quantized, f"{args.output_folder}/{Path(args.artifact).stem}_quantized.pt2"
    )


def compile(args):