oneDNN量化编程模型详解：INT8推理的完整指南-优快云博客

本文链接：https://blog.youkuaiyun.com/gitblog_00484/article/details/148552651

oneDNN量化编程模型详解：INT8推理的完整指南

oneDNN oneAPI Deep Neural Network Library (oneDNN) 项目地址: https://gitcode.com/gh_mirrors/on/oneDNN

引言

在现代深度学习推理中，量化技术已成为优化模型性能的关键手段。oneDNN作为高性能深度学习计算库，提供了完善的量化支持，特别是对INT8数据类型的优化。本文将深入解析oneDNN中的量化编程模型，帮助开发者理解并正确使用其量化功能。

量化基础概念

量化是将浮点数值转换为低精度整数表示的过程，其核心公式为：

x_f32[:] = scale_x * (x_int8[:] - zp_x)

其中：

scale_x：浮点缩放因子
zp_x：int32类型的零点偏移
[:]：表示公式对数组的逐元素应用

oneDNN不负责计算这些量化参数，而是由用户提供。量化参数可以通过两种方式确定：

静态量化：使用校准工具预先计算
动态量化：运行时根据张量的实际最小/最大值计算

oneDNN量化实现原理

基本工作流程

创建图元描述符：当输入为INT8时，图元将作为量化整数运算执行
设置量化参数维度：通过掩码指定缩放因子和零点的维度（如每个张量一个缩放因子、每个通道一个缩放因子等）
执行图元：在执行时提供实际的量化参数（f32类型的缩放因子和s32类型的零点）

数值行为特点

数据类型转换：允许将INT8输入转换为更宽的数据类型（如int16或int32）
中间计算：使用更宽的数据类型（如int32）进行中间计算和累加，避免溢出
结果转换：在写入输出内存前进行饱和转换（如s32到int8）

量化卷积示例

考虑带偏置的卷积运算，各张量表示为：

输入：src_f32[:] = scale_src * (src_int8[:] - zp_src)
权重：weights_f32[:] = scale_weights * weights_int8[:]
输出：dst_f32[:] = scale_dst * (dst_int8[:] - zp_dst)

实际计算过程为：

dst_int8[:] = f32_to_int8(
    (scale_src * scale_weights * s32_to_f32(conv_s32(src_int8, weights_int8))
    - zp_src * comp_s32 + bias_f32) / scale_dst + zp_dst)

其中：

conv_s32：常规卷积，使用int8输入并产生int32结果
comp_s32：补偿项，由oneDNN库计算
f32_to_s8：带饱和的f32到int8转换
s32_to_f32：int32到f32转换

每通道量化

oneDNN支持对权重的每输出通道缩放，计算方式为：

dst_int8(n, oc, oh, ow) = 
    f32_to_int8(
        (scale_src * scale_weights(oc) * 
        conv_s32(src_int8, weights_int8)|(n,oc,oh,ow) + bias_f32)/scale_dst
    )

权重准备需要使用特殊的重排序操作：

weights_int8(oc, ic, kh, kw) = 
    f32_to_int8(weights_f32(oc, ic, kh, kw) / scale_weights(oc))

API使用详解

基本缩放设置

通过primitive_attr::set_scales_mask设置缩放属性，将运算行为从：

dst[:] = Op(...)

变为：

dst[:] = scale * Op(...)

维度掩码设置

对于D0×...×Dn-1维度的张量，要为di维度设置每维度缩放，则：

掩码：mask = Σ(2^di)
缩放因子数量：scales.size() = Π(Ddi)

代码示例

权重每输出通道量化

// 设置权重量化属性
dnnl::primitive_attr attr;
const int quantization_mask = 0 | (1 << 0);  // OC维度(维度0)每通道缩放
attr.set_scales_mask(DNNL_ARG_DST, quantization_mask);

// 创建执行量化的重排序图元
auto wei_reorder_pd = dnnl::reorder::primitive_desc(
        wei_plain_f32_md, engine, // 源
        wei_conv_s8_md, engine,    // 目标
        attr);

每输出通道量化卷积

// 设置卷积属性
dnnl::primitive_attr attr;
const int data_mask = 0;  // 张量级缩放和零点
const int wei_mask = 0 | (1 << 0);  // 权重OC维度每通道缩放

attr.set_scales_mask(DNNL_ARG_SRC, data_mask);
attr.set_zero_points_mask(DNNL_ARG_SRC, data_mask);
attr.set_scales_mask(DNNL_ARG_WEIGHTS, wei_mask);
attr.set_scales_mask(DNNL_ARG_DST, data_mask);
attr.set_zero_points_mask(DNNL_ARG_DST, data_mask);

// 创建卷积图元描述符
auto conv_pd = dnnl::convolution_forward::primitive_desc(
        dnnl::prop_kind::forward_inference,
        dnnl::algorithm::convolution_direct,
        src_conv_s8_any_md, wei_conv_s8_any_md, dst_conv_s8_any_md,
        strides, padding_l, padding_r,
        dnnl::padding_kind::zero,
        attr);