OpenPPL PPQ量化(2)：离线静态量化源码剖析

目录

模型支持

量化onnx原生模型quantize_onnx_model

输入输出

执行流程

ONNX格式解析

后记

模型支持

openppl支持了三种模型onnx、caffe、pytorch其中pytorch和caffe是通过quantize_torch_model和quantize_caffe_model先将模型转换成onnx模型再调用quantize_onnx_model来实现量化的。

@ empty_ppq_cache
def quantize_torch_model(
    model: torch.nn.Module,
    calib_dataloader: DataLoader,
    calib_steps: int,
    input_shape: List[int],
    platform: TargetPlatform,
    input_dtype: torch.dtype = torch.float,
    setting: QuantizationSetting = None,
    collate_fn: Callable = None,
    inputs: List[Any] = None,
    do_quantize: bool = True,
    onnx_export_file: str = 'onnx.model',
    device: str = 'cuda',
    verbose: int = 0,
    ) -> BaseGraph:
   
    # dump pytorch model to onnx
    dump_torch_to_onnx(model=model, onnx_export_file=onnx_export_file,
        input_shape=input_shape, input_dtype=input_dtype,
        inputs=inputs, device=device)

    return quantize_onnx_model(onnx_import_file=onnx_export_file,
        calib_dataloader=calib_dataloader, calib_steps=calib_steps, collate_fn=collate_fn,
        input_shape=input_shape, input_dtype=input_dtype, inputs=inputs, setting=setting,
        platform=platform, device=device, verbose=verbose, do_quantize=do_quantize)

@ empty_ppq_cache
def quantize_caffe_model(
    caffe_proto_file: str,
    caffe_model_file: str,
    calib_dataloader: DataLoader,
    calib_steps: int,
    input_shape: List[int],
    platform: TargetPlatform,
    input_dtype: torch.dtype = torch.float,
    setting: QuantizationSetting = None,
    collate_fn: Callable = None,
    inputs: List[Any] = None,
    do_quantize: bool = True,
    device: str = 'cuda',
    verbose: int = 0,
) -> BaseGraph:
    if do_quantize:
        if calib_dataloader is None or calib_steps is None:
            raise TypeError('Quantization needs a valid calib_dataloader and calib_steps setting.')

    if setting is None:
        setting = QuantizationSettingFactory.default_setting()

    ppq_ir = load_graph(file_path=caffe_proto_file,
                        caffemodel_path=caffe_model_file,
                        from_framework=NetworkFramework.CAFFE)

    ppq_ir = format_graph(ppq_ir)
    ppq_ir = dispatch_graph(ppq_ir, platform, 
                            dispatcher=setting.dispatcher, 
                            dispatching_table=setting.dispatching_table)

    if inputs is None:
        dummy_input = torch.zeros(size=input_shape, device=device, dtype=input_dtype)
    else: dummy_input = inputs

    quantizer = PFL.Quantizer(platform=platform, graph=ppq_ir)
    executor = TorchExecutor(graph=quantizer._graph, device=device)
    executor.tracing_operation_meta(inputs=dummy_input)

    if do_quantize:
        quantizer.quantize(
            inputs=dummy_input,
            calib_dataloader=calib_dataloader,
            executor=executor,
            setting=setting,
            calib_steps=calib_steps,
            collate_fn=collate_fn
        )
        if verbose: quantizer.report()
        return quantizer._graph
    else:
        return quantizer._graph

所以我们接下来看看quantize_onnx_model是怎么实现的。

量化onnx原生模型quantize_onnx_model

输入输出

        onnx_import_file (str): 被量化的 onnx 模型文件路径 onnx model location

        calib_dataloader (DataLoader): 校准数据集 calibration data loader

        calib_steps (int): 校准步数 calibration steps

        collate_fn (Callable): 校准数据的预处理函数 batch collate func for preprocessing

        input_shape (List[int]): 模型输入尺寸用于执行 jit.trace对于动态尺寸的模型输入一个模型可接受的尺寸即可。
            如果模型存在多个输入则需要使用 inputs 变量进行传参此项设置为 None
                                a list of ints indicating size of input, for multiple inputs, please use
                                keyword arg inputs for direct parameter passing and this should be set to None

        input_dtype (torch.dtype): 模型输入数据类型如果模型存在多个输入则需要使用 inputs 变量进行传参此项设置为 None
                                the torch datatype of input, for multiple inputs, please use keyword arg inputs
                                for direct parameter passing and this should be set to None

        inputs (List[Any], optional): 对于存在多个输入的模型在Inputs中直接指定一个输入List从而完成模型的tracing。
                                for multiple inputs, please give the specified inputs directly in the form of
                                a list of arrays

        setting (OptimSetting): 量化配置信息用于配置量化的各项参数设置为 None 时加载默认参数。
                                Quantization setting, default setting will be used when set None

        do_quantize (Bool, optional): 是否执行量化 whether to quantize the model, defaults to True.


        platform (TargetPlatform, optional): 量化的目标平台 target backend platform, defaults to TargetPlatform.DSP_INT8.

        device (str, optional): 量化过程的执行设备 execution device, defaults to 'cuda'.

        verbose (int, optional): 是否打印详细信息 whether to print details, defaults to 0.

执行流程

我们首先要加载计算图

    ppq_ir = load_onnx_graph(onnx_import_file=onnx_import_file)

此处加载的计算图是原始的尚未被调度也就是所有算子都被认为是可量化的。

然后我们需要执行图的切分与调度不同算子会被执行不同的调度

    ppq_ir = dispatch_graph(graph=ppq_ir, platform=platform, 
                            dispatcher=setting.dispatcher, 
                            dispatching_table=setting.dispatching_table)

所有对计算图执行的操作最后都会返回BaseGraph类这个类是PPQ内部专门为模型量化准备的计算图除了保存一般计算图的必要信息之外还存储了所有量化信息。后面在写博客解析这个量化计算图的设计。

然后根据指定的平台platform确定指定的量化类

    quantizer = PFL.Quantizer(platform, ppq_ir)

所有的平台类型写在ppq/lib/common.py文件中

这些具体量化方法写在quantizer文件夹中传入量化计算图是因为这些量化类需要计算图进行初始化

因为我们已经初始化了量化类所以后面表示计算图不再使用ppq_ir直接用quantizer._graph表示。我们继续要用量化图初始化执行引擎这个引擎TorchExecutor能执行onnx的推理由于不同平台的推理细节是不同的所以这里的实现有点复杂大致的流程如下

详细的解析后面专门再写博客讲吧。

好了继续回到我们的主逻辑中最后一步是执行量化返回量化后的量化计算图搞定~

    if do_quantize:
        quantizer.quantize(
            inputs=dummy_input,
            calib_dataloader=calib_dataloader,
            executor=executor,
            setting=setting,
            calib_steps=calib_steps,
            collate_fn=collate_fn
        )
        if verbose: quantizer.report()
        return quantizer._graph
    else:
        executor = TorchExecutor(graph=ppq_ir, device=device)
        executor.tracing_operation_meta(inputs=

最后注意这里如果不需要执行量化我们用没有原始载入的计算图执行一遍推理然后返回即可。

ONNX格式解析

如果不了解ONNX格式前面从ONNX解析出计算图部分会比较难理解有一篇写的很棒的博客我摘抄了一部分帮助理解ONNX学习笔记 - 知乎

这一节我们来分析一下ONNX的组织格式上面提到ONNX中最核心的部分就是onnx.protohttps://github.com/onnx/onnx/blob/master/onnx/onnx.proto这个文件了它定义了ONNX这个数据协议的规则和一些其它信息。现在是2021年1月这个文件有700多行我们没有必要把这个文件里面的每一行都贴出来我们只要搞清楚里面的核心部分即可。在这个文件里面以message关键字开头的对象是我们需要关心的。我们列一下最核心的几个对象并解释一下它们之间的关系。

• ModelProto
• GraphProto
• NodeProto
• ValueInfoProto
• TensorProto
• AttributeProto

当我们加载了一个ONNX之后我们获得的就是一个ModelProto它包含了一些版本信息生产者信息和一个GraphProto。在GraphProto里面又包含了四个repeated数组它们分别是node(NodeProto类型)input(ValueInfoProto类型)output(ValueInfoProto类型)和initializer(TensorProto类型)其中node中存放了模型中所有的计算节点input存放了模型的输入节点output存放了模型中所有的输出节点initializer存放了模型的所有权重参数。

我们知道要完整的表达一个神经网络不仅仅要知道网络的各个节点信息还要知道它们的拓扑关系。这个拓扑关系在ONNX中是如何表示的呢ONNX的每个计算节点都会有input和output两个数组这两个数组是string类型通过input和output的指向关系我们就可以利用上述信息快速构建出一个深度学习模型的拓扑图。这里要注意一下GraphProto中的input数组不仅包含我们一般理解中的图片输入的那个节点还包含了模型中所有的权重。例如Conv层里面的W权重实体是保存在initializer中的那么相应的会有一个同名的输入在input中其背后的逻辑应该是把权重也看成模型的输入并通过initializer中的权重实体来对这个输入做初始化即一个赋值的过程。

最后每个计算节点中还包含了一个AttributeProto数组用来描述该节点的属性比如Conv节点或者说卷积层的属性包含grouppadstrides等等每一个计算节点的属性输入输出信息都详细记录在https://github.com/onnx/onnx/blob/master/docs/Operators.md。

后记

关于如何做量化校准如果使用校准数据如何配置量化设置具体的量化过程是如何如何选择需要量化的算子……

还有很多问题没有讲明白这个系列很长我们一一探索