[人工智能]ONNX的模型优化与量化细节

layout: post
title: ONNX的模型优化与量化细节
date: 2021-09-21 18:18:48.000000000 +09:00
categories: [算法框架]
tags: [离线推理]

ONNX的模型优化与量化细节

ONNX基本介绍

• 什么是ONNX？

  ONNX全称为 Open Neural Network Exchange，是一种与框架无关的模型表达式。ONNX的规范及代码主要由微软，亚马逊 ，Facebook 和 IBM 等公司共同开发，以开放源代码的方式托管在Github上。目前官方支持加载ONNX模型并进行推理的深度学习框架有： Caffe2, PyTorch, MXNet，ML.NET，TensorRT 和 Microsoft CNTK，并且 TensorFlow 也非官方的支持ONNX。

• ONNX的数据格式是怎么样的？

  ONNX本质上一种文件格式，通过Protobuf数据结构存储了神经网络结构权重。其组织格式核心定义在于onnx.proto，其中定义了Model/Graph/Node/ValueInfo/Tensor/Attribute层面的数据结构。整图通过各节点（Node）的input/output指向关系构建模型图的拓扑结构。

• ONNX支持的功能？

  基于ONNX模型，官方提供了一系列相关工具：模型转化/模型优化（simplifier等）/模型部署(Runtime)/模型可视化（Netron等）等。

ONNX自带了Runtime库，能够将ONNX Model部署到不同的硬件设备上进行推理，支持各种后端（如TensorRT/OpenVINO）。

基于ONNX Model的Runtime系统架构如下，可以看到Runtime实现功能是将ONNX Model转换为In-Memory Graph格式，之后通过将其转化为各个可执行的子图，最后通过GetCapability() API将子图分配到不同的后端（execution provider）执行。

ONNX模型优化

onnx_simplifier的核心功能如下：

ONNX Simplifier is presented to simplify the ONNX model. It infers the whole computation graph and then replaces the redundant operators with their constant outputs.

simplify的基本流程如下：

1. 利用onnxruntime推理计算图，得到各个节点的输入输出的infer shape
2. 基于ONNX支持的优化方法进行ONNX模型的优化（如fuse_bn_into_conv）
3. 对ONNX模型的常量OP进行折叠：
   1. 基于get_constant_nodes获取常量OP
   2. 基于add_features_to_output将所有静态节点的输出扩展到ONNX图的输出节点列表中（主要为了后续步骤方便获取常量节点输出）
   3. 将1.中得到的常量OP从图中移除（断开连线），同时将其节点参数构建为其他节点的输入参数
   4. 清理图中的孤立节点（3.中断开连线的节点）

其optimize函数定义如下（可以简单看下目前支持的一些优化方法）：

def optimize(model: onnx.ModelProto, skip_fuse_bn: bool, skipped_optimizers: Optional[Sequence[str]]) -> onnx.ModelProto:
    """
    :model参数: 待优化的ONXX模型.
    :return: 优化之后的ONNX模型.
    简化之前, 使用这个方法产生会在'forward_all'用到的ValueInfo
    简化之后，使用这个方法去折叠前一步产生的常量到initializer中并且消除没被使用的常量
    """
    onnx.checker.check_model(model)
    onnx.helper.strip_doc_string(model)
    optimizers_list = [
        'eliminate_deadend',
        'eliminate_nop_dropout',
        'eliminate_nop_cast',
        'eliminate_nop_monotone_argmax', 'eliminate_nop_pad',
        'extract_constant_to_initializer', 'eliminate_unused_initializer',
        'eliminate_nop_transpose',
        'eliminate_nop_flatten', 'eliminate_identity',
        'fuse_add_bias_into_conv',
        'fuse_consecutive_concats',
        'fuse_consecutive_log_softmax',
        'fuse_consecutive_reduce_unsqueeze', 'fuse_consecutive_squeezes',
        'fuse_consecutive_transposes', 'fuse_matmul_add_bias_into_gemm',
        'fuse_pad_into_conv', 'fuse_transpose_into_gemm', 'eliminate_duplicate_initializer'
    ]
    if not skip_fuse_bn:
        optimizers_list.append('fuse_bn_into_conv')
    if skipped_optimizers is not None:
        for opt in skipped_optimizers:
            try:
                optimizers_list.remove(opt)
            except ValueError:
                pass

    model = onnxoptimizer.optimize(model, optimizers_list, fixed_point=True)
    onnx.checker.check_model(model)
    return model

ONNX转FP16

ONNX支持FP32模型转换为FP16模型，接口如下：

import onnxmltools
from onnxmltools.utils.float16_converter import convert_float_to_float16

# Update the input name and path for your ONNX model
input_onnx_model = 'model.onnx'
# Change this path to the output name and path for your float16 ONNX model
output_onnx_model = 'model_f16.onnx'
# Load your model
onnx_model = onnxmltools.utils.load_model(input_onnx_model)
# Convert tensor float type from your input ONNX model to tensor float16
onnx_model = convert_float_to_float16(onnx_model)
# Save as protobuf
onnxmltools.utils.save_model(onnx_model, output_onnx_model)

具体实现在于float16.py，截断的逻辑：

• 小于最小精度（默认1e-7）映射为最小精度
• 大于最大范围（默认1e4）映射为最大值
• NaN/0/inf/-inf保持原值

核心代码：

def convert_np_to_float16(np_array, min_positive_val=1e-7, max_finite_val=1e4):
    def between(a, b, c):
        return np.logical_and(a < b, b < c)
    np_array = np.where(between(0, np_array, min_positive_val), min_positive_val, np_array)
    np_array = np.where(between(-min_positive_val, np_array, 0), -min_positive_val, np_array)
    np_array = np.where(between(max_finite_val, np_array, float('inf')), max_finite_val, np_array)
    np_array = np.where(between(float('-inf'), np_array, -max_finite_val), -max_finite_val, np_array)
    return np.float16(np_array)

参考资料

onnxruntime官方资料

onnx比较好的知乎介绍

onnx一些调研资料

onnx_simplifier常量折叠介绍