深入理解oneDNN中的后操作(Post-ops)属性-优快云博客

本文链接：https://blog.youkuaiyun.com/gitblog_00959/article/details/148552645

深入理解oneDNN中的后操作(Post-ops)属性

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

后操作概述

在深度学习计算中，操作融合(Operation Fusion)是一种重要的性能优化技术。oneDNN通过后操作属性API实现了基本的操作融合能力。后操作(Post-ops)是指在主操作(如卷积、矩阵乘法等)之后追加执行的一系列操作，通过减少内存带宽压力来提高性能。

后操作使用属性机制实现，如果存在多个后操作，它们将按照添加顺序依次执行。目前oneDNN支持以下几种后操作类型：

元素级操作(Eltwise)
求和操作(Sum)
深度卷积(Depthwise)
二元操作(Binary)
PReLU操作

后操作的基本使用

后操作通过不透明的结构体表示(在C API中是dnnl_post_ops_t，在C++ API中是dnnl::post_ops)。下面是一个基本的使用示例：

dnnl::post_ops po; // 创建空的后操作对象
po.append_SOMETHING(params); // 添加第一个后操作
po.append_SOMETHING_ELSE(other_params); // 添加第二个后操作

dnnl::primitive_attr attr; // 创建属性对象
attr.set_post_ops(po); // 将后操作附加到属性上

// 使用带有后操作属性的描述符创建原语
primitive::primitive_desc op_pd(engine, params, attr);

重要说明：

后操作的添加顺序决定了它们的执行顺序
库支持的最大后操作数量为32
不同原语对后操作的支持程度不同，需要查阅具体原语的文档
后操作不会改变目标内存对象的格式

支持的后操作类型详解

1. 元素级后操作(Eltwise Post-op)

元素级后操作可以将主操作(如卷积或内积)与元素级操作(通常是激活函数)融合。

API:

void append_eltwise(algorithm alg, float alpha, float beta);

数学表示：原始操作：dst = Op(...) 添加后操作后变为：dst = eltwise(Op(...))

典型应用：在卷积后直接应用ReLU等激活函数。

注意：中间结果Op(...)不会被保留，因此这种融合通常不能用于训练阶段。

2. 求和后操作(Sum Post-op)

求和后操作将主操作的结果与现有数据累加，累加前可以对现有数据进行缩放和偏移。

数学表示：原始操作：dst = Op(...) 添加后操作后变为：dst = scale * (dst - zero_point) + Op(...)

典型应用：残差学习块，其中卷积结果需要与先前计算的激活值相加。

数据类型支持：

可以指定目标张量的数据类型进行重新解释
数据类型大小必须与原始目标数据类型相同

3. 深度卷积后操作(Depthwise Post-op)

深度卷积后操作将1x1卷积与深度卷积融合，这在MobileNet等模型架构中很常见。

数学表示：原始操作：dst = Conv_1x1(...) 添加后操作后变为：dst = Conv_dw(Conv_1x1(...))

输出维度：最终输出维度为：{n, oc_1x1, ceil(oh_conv_1x1/stride), ceil(ow_conv_1x1/stride)}

支持的数据类型组合：

| 1x1卷积输出类型 | 深度卷积输出类型 | 深度卷积权重类型 | 深度卷积偏置类型 | |----------------|-----------------|----------------|----------------| | u8, s8 | u8, s8, s32, f32| s8 | f32, s32 | | f32 | f32 | f32 | f32 | | bf16 | bf16, f32 | bf16 | f32, bf16 | | f16 | f16, f32 | f16 | f32, f16 |

重要限制：

目前仅支持2D 1x1卷积
不能与其他深度卷积或求和后操作链式组合
目标内存格式必须为any
当深度卷积步长≠1时，需要特别注意空间维度处理

4. 二元后操作(Binary Post-op)

二元后操作可以将主操作与二元操作融合。

API:

void append_binary(algorithm alg, const memory::desc &src1);

数学表示：原始操作：dst = Op(...) 添加后操作后变为：dst = binary(Op(...), Source_1[:])

优化场景：

张量级广播(Source_1为单元素张量)
通道级广播(Source_1的通道数与主操作相同)
元素级广播(Source_1与主操作完全匹配)

选择操作(Select)：对于二元选择操作，需要额外的条件张量：

void append_binary(algorithm alg, 
                  const memory::desc &src1,
                  const memory::desc &src2);

5. PReLU后操作(PReLU Post-op)

PReLU后操作可以将主操作与PReLU操作融合。

API:

void append_prelu(int mask);

数学表示：原始操作：dst = Op(...) 添加后操作后变为：dst = prelu(Op(...), weights[:])

关键参数：

mask：描述权重广播方式的掩码
权重张量在运行时通过特定机制传递
目前仅支持fp32权重和plain布局

后操作链示例

示例1：求和后接ReLU

这是ResNet家族网络中的常见模式。

dnnl::post_ops po;
po.append_sum();
po.append_eltwise(dnnl::algorithm::eltwise_relu, 0.f, 0.f);

dnnl::primitive_attr attr;
attr.set_post_ops(po);

数学表示：dst = ReLU(dst + conv(src, weights))

示例2：Tanh -> 求和 -> 线性变换

这个例子展示了操作序列和缩放的使用。

dnnl::post_ops po;
po.append_eltwise(dnnl::algorithm::eltwise_tanh, 0.f, 0.f);
po.append_sum();
po.append_eltwise(dnnl::algorithm::eltwise_linear, alpha, beta);

dnnl::primitive_attr attr;
attr.set_scales_mask(DNNL_ARG_SRC, 0);
attr.set_scales_mask(DNNL_ARG_WEIGHTS, 0);
attr.set_scales_mask(DNNL_ARG_DST, 0);
attr.set_post_ops(po);

数学表示：dst = s_linear * (α * (s_sum * dst + s_tanh * tanh(s_conv * conv(src, weights))) + β)

示例3：ReLU -> 深度卷积 -> ReLU

这是MobileNet中融合深度卷积和1x1卷积的示例。

dnnl::post_ops po;
po.append_eltwise(dnnl::algorithm::eltwise_relu, 0.f, 0.f);
po.append_dw(dnnl::memory::data_type::s8, 
             dnnl::memory::data_type::undef,
             dnnl::memory::data_type::u8,
             kernel, stride, padding);
po.append_eltwise(dnnl::algorithm::eltwise_relu, 0.f, 0.f);

dnnl::primitive_attr attr;
attr.set_scales_mask(DNNL_ARG_DST, 0);
attr.set_scales_mask(DNNL_ARG_ATTR_POST_OP_DW | DNNL_ARG_DST, 0);
attr.set_post_ops(po);

数学表示：dst = ReLU_depthwise(scales_depthwise * (conv_depthwise(ReLU_1x1(scales_conv_1x1 * (conv_1x1()))))