X5平台PTQ权重拆分调优说明

背景

对于X5计算平台,卷积算子在正常情况下,最高只能支持到W8A16量化精度,即权重int8,激活int16。但对于复杂模型,W8A16越来越难满足精度要求,因此本文提供一种名为权重拆分的调优方法,可以变相实现卷积的W16A16量化,以提升模型的精度表现。

流程详解

  1. 按正常PTQ流程找到最佳的yaml参数组合

本文介绍的调优方法是正常调优流程无法满足时的补充,因此需要用户先使用常规手段找到模型精度表现最好的yaml组合,如max,max_percentile,per_channel等等,并且set_all_nodes_int16通常都需要打开。

  1. 做weight的精度debug

按照手册的精度debug指南章节,进行weight的精度debug。

# 导入debug模块
import horizon_nn.quantizer.debugger as dbg
# 导入log日志模块
import logging

# 若verbose=True时,需要先设置log level为INFO
logging.getLogger().setLevel(logging.INFO)
# 获取节点量化敏感度
node_message = dbg.get_sensitivity_of_nodes(
        model_or_file='./calibrated_model.onnx',
        metrics=['cosine-similarity', 'mse'],
        calibrated_data='./calibration_data/',
        output_node=None,
        node_type='weight',
        data_num=None,
        verbose=True,
        interested_nodes=None)

用户可使用nohup python3 debug.py >debug_output.log 2>&1 &等方式让精度debug结果保存到本地。

之后找到精度debug结果中,相似度低于0.9999的conv算子,记录下完整的node name。

  1. 权重拆分

权重拆分是指利用特定方法,让原本不支持conv weight int16的工具链变相支持这个功能,可进一步提升精度,对于conv weight int8精度不够导致的掉点问题尤其有效。

生成不带预处理节点的optimized.onnx

做权重拆分的前提是,我们需要先得到一个PTQ生成的optimized.onnx模型,并且这个模型不能有预处理节点,因此可以让模型以featuremap作为输入来修改yaml文件。在这一步,我们得到optimized.onnx即可终止程序。

input_parameters:
  input_type_rt: 'featuremap'
  input_layout_rt: 'NCHW'
  input_type_train: 'featuremap'
  input_layout_train: 'NCHW'
  input_shape: '1x3x256x128'

基于optimized.onnx做权重拆分

权重拆分的python脚本如下,脚本的输入为上一步得到的optimized.onnx模型,conv_names使用精度debug中得到的余弦相似度<0.9999的node,最终生成xxx_split.onnx模型。

import numpy as np  
from copy import deepcopy

# horizon_nn 1.1.0
from horizon_nn.common import constant_folding
from horizon_nn.ir import load_model, save_model

# horizon_nn newer version
# from hmct.common import constant_folding
# from hmct.it import load_model, save_model

# 最近的patch也有可能改成了这个
# from hmct.common import ConstantFolding
# from hmct.ir import load_model, save_model

def split_conv_nodes(model, conv_names):
    for conv_name in conv_names:
        conv_node = model.graph.node_mappings[conv_name]
        before_node = conv_node.inputs[0].src_op
        conv_weight_value = deepcopy(conv_node.inputs[1].value)
        conv_weight_max = abs(conv_weight_value).max(axis=(1,2,3))
        moded = (conv_weight_max / 127)[:, np.newaxis, np.newaxis, np.newaxis] + 1e-10
        conv_weight_high = np.floor(np.clip(conv_weight_value / moded + 1e-5, -127, 127)) * moded
        conv_weight_low = conv_weight_value - conv_weight_high
        conv_bias_value = conv_node.inputs[2].value if len(conv_node.inputs) == 3 else np.zeros(conv_weight_value.shape[0], np.float32)
        conv1_weight_var = model.graph.create_variable(
            is_param=True,
            value=conv_weight_high,
        )
        conv1_bias_var = conv_node.inputs[2] if len(conv_node.inputs) == 3 else model.graph.create_variable(
            is_param=True,
            value=np.zeros_like(conv_bias_value, np.float32),
        )
        conv1_node = model.graph.create_node(
            op_type="Conv",
            name = conv_node.name + "_split0",
            attributes=conv_node.attributes,
            inputs=[conv_node.inputs[0], conv1_weight_var, conv1_bias_var],
            num_outputs=1)
        if before_node is not None:
            conv1_node.insert_after(before_node)
        else:
            conv1_node.prepend_on()
        conv2_weight_var = model.graph.create_variable(
            is_param=True,
            value=conv_weight_low,
        )
        conv2_bias_var = model.graph.create_variable(
            is_param=True,
            value=np.zeros_like(conv_bias_value, np.float32),
        )
        conv2_node = model.graph.create_node(
            op_type="Conv",
            name = conv_node.name + "_split1",
            attributes=conv_node.attributes,
            inputs=[conv_node.inputs[0], conv2_weight_var, conv2_bias_var],
            num_outputs=1)
        if before_node is not None:
            conv2_node.insert_after(before_node)
        else:
            conv2_node.prepend_on()
        add1_node = model.graph.create_node(
            op_type="Add",
            inputs=[conv1_node.outputs[0], conv2_node.outputs[0]],
            name=conv_node.name + "_split_add0",
            num_outputs=1).insert_after(conv1_node)
        conv_node.replace_all_uses_with(add1_node)
        if not conv_node.is_used:
            conv_node.destroy()
    model.infer_shapes()
    model.check_validity()
    return model

if __name__ == "__main__":
    model = constant_folding(load_model("xxx_optimized_float_model.onnx"))
    model = split_conv_nodes(model, conv_names=[
        "/backbone/mod1/mod1.0/Conv",
        "/backbone/mod1/mod1.0_1/Conv",
        "/backbone/mod2/mod2.0/head_layer/conv/conv.0/conv.0.0/Conv",
    ])
    save_model(model, "xxx_split.onnx")

脚本读取的模型为第一步得到的xxx_optimized_float_model.onnx文件,split_conv_nodes的参数使用第二步得到的weight余弦相似度低于0.9999的卷积的node name,最后使用save_model将权重拆分后的onnx模型保存。

  1. 基于split onnx再走一遍PTQ

使用保存的xxx_split.onnx再走一遍PTQ流程,正常情况下我们可以得到精度更高的量化模型。

可视化对比

以下为权重拆分前后的量化模型检测结果对比。

2 个赞