ARTICLE DETAIL

资讯详情

深耕网站视觉设计与运营推广的一线实战洞察。

CANN ops-nn GroupNorm 算子深度解析:NPU 分组归一化的接口、约束与源码实现

CANN ops-nn GroupNorm 算子深度解析:NPU 分组归一化的接口、约束与源码实现 CANN ops-nn GroupNorm 算子深度解析NPU 分组归一化的接口、约束与源码实现【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nnGroupNorm分组归一化是深度学习中重要的归一化算子在 batch size 较小或训练/推理场景中常作为 BatchNorm 的替代方案。本文以 CANN ops-nn 算子仓库中 experimental/norm/group_norm/README.md 为核心结合该模块的 op_api、op_host、op_kernel 与单元测试源码系统讲解 GroupNorm 算子的数学原理、产品支持情况、参数与约束、aclnn 调用方式以及 tiling 调度与 AI Core Kernel 的底层实现帮助读者完整掌握在 NPU 上使用与理解该算子所需的知识。一、功能说明与数学原理根据 README 功能说明该算子对输入张量self按group维度进行归一化计算并可选使用gamma和beta对归一化结果做仿射变换将结果写入out同时输出每个样本每个 group 的均值meanOut和倒标准差rstdOut。归一化按「样本 × 分组」为单位进行设某个样本n的某个分组g内元素个数为m即channelsPerGroup × HxW其计算过程分为三步计算分组均值$$ mean_{n,g} \frac{1}{m}\sum_{i \in group(g)} x_i $$计算分组倒标准差rstd即 1/标准差同时引入数值稳定项eps$$ rstd_{n,g} \frac{1}{\sqrt{\frac{1}{m}\sum_{i \in group(g)}(x_i - mean_{n,g})^2 eps}} $$归一化与仿射变换gamma、beta均按通道 c 取标量值同一个 group 内的所有通道共享该 group 的 mean 与 rstd$$ y_i (x_i - mean_{n,g}) \times rstd_{n,g} \times gamma_c beta_c $$当gamma或beta为空时分别按gamma 1、beta 0处理。这一约定在 Kernel 实现中也有直接印证在 group_norm_kernel.h 的NormalizeFloatTile中tilingData.hasGamma 0时取gammaVal 1.0ftilingData.hasBeta 0时取betaVal 0.0f。从语义上看该算子的输出同时包含归一化结果out与统计量meanOut/rstdOut这在需要在前向传播后继续做反向传播如 LayerNorm 类算子复用统计量的场景时非常实用无需重新扫描输入即可计算梯度。二、产品支持情况根据 README 产品支持情况该算子在当前仓库版本下的支持矩阵如下产品是否支持Ascend 950PR / Ascend 950DT×Atlas A3 训练系列产品 / Atlas A3 推理系列产品√Atlas A2 训练系列产品 / Atlas A2 推理系列产品√Atlas 200I/500 A2 推理产品×Atlas 推理系列产品×Atlas 训练系列产品×同时从算子定义源码 group_norm_def.cpp 可以看到当前实现的 AI Core 配置仅注册了ascend910b对应 Atlas A2/A3 训练与推理系列所基于的架构平台这与 README 的支持矩阵一致。需要特别说明该算子位于experimental/目录下属于实验性算子使用前请以当前发布版本的官方支持清单为准。三、参数说明README 给出了完整的参数表。该算子共有 8 个输入参数与 3 个输出参数其中N、C、HxW、group、eps为标量属性参数。汇总如下参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续 Tensorself输入GroupNorm 的输入张量公式中的 xshape 需满足 rank 为 2-8且第 0 维等于 N、第 1 维等于 C第 2 维及之后元素个数乘积等于 HxWFLOAT、FLOAT16、BFLOAT16ND2-8√gamma可选输入仿射变换的缩放系数公式中的 gamma可为空不为空时 shape 需为 [C]数据类型需与 self 一致FLOAT、FLOAT16、BFLOAT16ND1√beta可选输入仿射变换的偏置系数公式中的 beta可为空不为空时 shape 需为 [C]数据类型需与 self 一致FLOAT、FLOAT16、BFLOAT16ND1√N输入输入张量的 batch 维大小需等于 self.shape[0]INT64---C输入输入张量的 channel 维大小需等于 self.shape[1]且能被 group 整除INT64---HxW输入输入张量从第 2 维开始的元素个数乘积需等于 self.shape[2] * ... * self.shape[rank-1]INT64---group输入GroupNorm 的分组数需大于 0且 C % group 0INT64---eps输入归一化计算中加到方差上的数值稳定项需大于 0DOUBLE---out输出归一化并仿射变换后的输出shape 和数据类型需与 self 一致FLOAT、FLOAT16、BFLOAT16ND2-8√meanOut输出每个样本每个 group 的均值shape 需为 [N, group]数据类型需与 self 一致FLOAT、FLOAT16、BFLOAT16ND2√rstdOut输出每个样本每个 group 的倒标准差shape 需为 [N, group]数据类型需与 self 一致FLOAT、FLOAT16、BFLOAT16ND2√几点需要特别注意数据类型输入输出张量仅支持FLOATfloat32、FLOAT16、BFLOAT16三种eps使用DOUBLEdouble类型其余标量参数使用INT64。该约束在算子定义 group_norm_def.cpp 中通过.DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16})注册。非连续 Tensor表中标注√的张量参数均支持非连续strided输入对视图view类张量友好。meanOut / rstdOut 的 shape 恒为 [N, group]这一点在 group_norm_infershape.cpp 中有严格实现——输出y的 shape 直接复制输入x而mean/rstd的 shape 被设置为 2 维[N, num_groups]。eps 默认值在算子定义中eps属性被声明为OPTIONAL且默认值为0.00001fgroup_norm_tiling.cpp 中同样处理了epsAttr nullptr时取0.00001f的情况。四、约束说明结合 README 与源码使用该算子需满足以下约束self的 rank 需为 2-8且N、C、HxW需与self的实际 shape 一致N self.shape[0]、C self.shape[1]、HxW self.shape[2]*...*self.shape[rank-1]。group需大于 0且C需能被group整除。eps需大于 0。gamma、beta为可选输入不为空时 shape 均需为[C]数据类型需与self一致。out、meanOut、rstdOut为必选输出数据类型需与self一致out的 shape 需与self一致meanOut和rstdOut的 shape 需为[N, group]。op_host tiling 路径不支持 N、C 或 HxW 为 0 的输入。这一点在 group_norm_tiling.cpp 中有显式校验OP_CHECK_IF(c 0 || n 0 || hxw 0, ...)会直接返回失败同时该处还校验了c % numGroups ! 0。从算子定义看该算子开启了DynamicCompileStaticFlag(true)、DynamicRankSupportFlag(true)、DynamicShapeSupportFlag(true)与PrecisionReduceFlag(true)即支持动态 shape 与动态 rank并允许在精度允许的情况下启用降低精度的优化。五、调用说明与 aclnn 调用样例README 中给出的调用方式为aclnn 调用即通过aclnnGroupNormGetWorkspaceSize与aclnnGroupNorm两步式接口完成算子执行。这两个接口在 aclnn_group_norm_experimental.h 中声明aclnnStatus aclnnGroupNormGetWorkspaceSize(const aclTensor* self, const aclTensor* gammaOptional, const aclTensor* betaOptional, int64_t n, int64_t c, int64_t hxw, int64_t group, double eps, aclTensor* out, aclTensor* meanOut, aclTensor* rstdOut, uint64_t* workspaceSize, aclOpExecutor** executor); aclnnStatus aclnnGroupNorm(void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, aclrtStream stream);使用流程与 CANN 标准 aclnn 接口一致先调用aclnnGroupNormGetWorkspaceSize获取 workspace 大小并创建执行器再调用aclnnGroupNorm在指定 stream 上异步执行。仓库中的单元测试 test_aclnn_group_norm.cpp 给出了可直接参照的调用样例。以最常见的 float16、四维输入、带 gamma/beta 的用例为例// 输入 shape [2, 8, 4, 4]即 N2, C8, HxW16group4 auto self TensorDesc({2, 8, 4, 4}, ACL_FLOAT16, ACL_FORMAT_ND); auto gamma TensorDesc({8}, ACL_FLOAT16, ACL_FORMAT_ND); auto beta TensorDesc({8}, ACL_FLOAT16, ACL_FORMAT_ND); int64_t n 2, c 8, hxw 16, group 4; double eps 0.00001; auto out TensorDesc({2, 8, 4, 4}, ACL_FLOAT16, ACL_FORMAT_ND); auto mean TensorDesc({2, 4}, ACL_FLOAT16, ACL_FORMAT_ND); // [N, group] auto rstd TensorDesc({2, 4}, ACL_FLOAT16, ACL_FORMAT_ND); // [N, group] auto ut OP_API_UT(aclnnGroupNorm, INPUT(self, gamma, beta, n, c, hxw, group, eps), OUTPUT(out, mean, rstd)); uint64_t workspaceSize 0; EXPECT_EQ(ut.TestGetWorkspaceSize(workspaceSize), ACLNN_SUCCESS);测试用例还覆盖了若干关键边界与异常场景可作为正确性参照无 gamma/beta 的可选输入group_norm_normal_float32_without_gamma_beta用例将gamma、beta均传nullptr验证可选输入为空时的路径此时按 gamma1、beta0 处理。空指针校验group_norm_null_self用例传入nullptr的self期望返回ACLNN_ERR_PARAM_NULLPTR。非法数据类型group_norm_invalid_dtype用例传入ACL_DOUBLE类型的张量期望返回ACLNN_ERR_PARAM_INVALID。非法分组数group_norm_invalid_group用例取group38 不能被 3 整除期望返回ACLNN_ERR_PARAM_INVALID。非法输出 shapegroup_norm_invalid_mean_shape用例将mean的 shape 设为[2, 3]与 group4 不符期望返回ACLNN_ERR_PARAM_INVALID。六、算子源码架构解析与 ops-nn 中其他算子的组织方式一致GroupNorm 模块按四层结构组织目录 experimental/norm/group_normexperimental/norm/group_norm/ ├── op_api/ # 对外 aclnn 接口与 L0 算子封装 ├── op_graph/ # 算子原型proto声明 ├── op_host/ # 算子定义、InferShape、Tiling宿主侧 ├── op_kernel/ # AI Core Kernel 实现与 tiling 数据结构 └── tests/ut/ # op_api 与 op_host 的单元测试各层职责与关键文件如下。6.1 op_api对外接口层aclnn_group_norm_experimental.h对外导出的aclnnGroupNormGetWorkspaceSize/aclnnGroupNorm两步接口声明。group_norm.cppL0 层实现。l0op::GroupNorm内部自动为y、mean、rstd分配输出张量——其中y复用输入x的 view shape 与数据类型mean/rstd则按Shape({n, numGroups})与FORMAT_ND分配然后通过ADD_TO_LAUNCHER_LIST_AICORE将算子加入 AI Core 启动列表。group_norm_l0.h 与 aclnn_group_norm.cppL0 算子类型注册与 aclnn 参数适配。6.2 op_graph算子原型group_norm_proto.h 声明算子的图级原型供 GE图引擎侧的 InferShape 测试与图编译使用。6.3 op_host算子定义、InferShape 与 Tiling算子定义group_norm_def.cpp输入x、gamma、betagamma、beta为OPTIONAL三者数据类型均为DT_FLOAT16 / DT_FLOAT / DT_BF16、格式FORMAT_ND并设置AutoContiguous()对非连续输入自动处理为连续。输出y、mean、rstd均为REQUIRED数据类型与格式同上。属性num_groupsREQUIREDInt、epsOPTIONALFloat默认0.00001f。AI Core 配置注册平台ascend910b开启动态 shape、动态 rank、动态编译支持。InferShapegroup_norm_infershape.cpp校验输入 rank 2、num_groups 0输出y的 shape 直接拷贝输入x输出mean、rstd的 shape 设为[N, num_groups]mean的第 0 维取xShape-GetDim(0)第 1 维取num_groupsrstd与mean相同。对应的单元测试 test_group_norm_infershape.cpp 验证了输入[2,8,4,4]、num_groups4时y保持[2,8,4,4]mean/rstd为[2,4]同时覆盖了 rank 非法如 1 维输入等失败场景。Tilinggroup_norm_tiling.cppTiling 是宿主侧为 Kernel 计算做「切分与调度」的关键环节GroupNorm 的 Tiling 逻辑清晰且值得展开读取 shape 与属性从输入 storage shape 读取N、C并累乘第 2 维及之后得到hxw随后计算channelsPerGroup C / numGroups、elementsPerGroup channelsPerGroup * hxw、invElementsPerGroup 1 / elementsPerGroup、groupNum N * numGroups。同时根据gamma/beta是否为空设置hasGamma/hasBeta标志。Core 分配blockNum min(groupNum, coreNum)每个 Core 平均分配groupPerCore个 group前groupTailCore个 Core 各多分 1 个 group即采用「均分 余数前移」的负载均衡策略。Tile 切分tileLength依据 UBUnified Buffer容量计算tileLength (ubSize / BUFFER_NUM) / sizeof(float) / 8下限为MIN_TILE_BYTES / sizeof(float)即 8KB并按 32 字节BLOCK_SIZE对齐若elementsPerGroup小于tileLength则直接取elementsPerGroup。可见双缓冲BUFFER_NUM 2是数据搬运的基本策略。Tiling Key 选择当elementsPerGroup tileLength时选择GROUP_NORM_SCH_SMALL_GROUP否则选择GROUP_NORM_SCH_GENERAL定义见 group_norm_tiling_key.h并将blockNum设为 block dim。Workspace通过platform.GetLibApiWorkSpaceSize()申请算子库所需的 workspace。6.4 op_kernelAI Core Kernel 实现Kernel 入口 group_norm.cpp 注册 tiling 数据后调用GroupNormKernel::GroupNormDTYPE_X, schMode::Process()核心实现在 group_norm_kernel.h。其执行流程可以归纳为「两趟扫描」统计趟CalcMeanAndVariance对每个 group 分 tile 搬运数据采用Kahan 补偿求和代码中的sumComp/squareSumComp即补偿项计算sum与squareSum从而得到高精度的mean与方差最终计算rstd 1 / sqrt(variance eps)。对 float16/bfloat16 输入会先Cast到 float 再统计保证中间精度。归一化输出趟NormalizeAndCopyOut再次分 tile 读取数据按通道分段NormalizeFloatTile中以hxw为粒度定位通道执行Adds(-mean)、Muls(rstd * gamma)、Adds(beta)对非 float 类型先转 float 计算再Cast回原类型RoundMode::CAST_ROUND写出到y同时将meanOut、rstdOut通过StoreScalar逐标量写回。此外 Kernel 还针对elementsPerGroup 1的退化场景每个 group 只有一个元素实现了专门的ProcessSingleElementGroups快速路径此时rstd恒为1/sqrt(eps)、归一化后输出即beta无 gamma 时避免了冗余的统计计算。tiling 数据由 group_norm_tiling_data.h 中的GroupNormTilingData结构体承载含n、c、hxw、numGroups、channelsPerGroup、elementsPerGroup、groupNum、groupPerCore、groupTailCore、tileLength、hasGamma、hasBeta、eps、invElementsPerGroup等字段该结构体由宿主侧 Tiling 填充、Kernel 侧读取是 host 与 device 通信的纽带。七、总结GroupNorm 是 CANN ops-nn 在experimental/norm下提供的实验性归一化算子支持 rank 2-8 的 ND 格式张量、float16/float/bfloat16 三种数据类型输出归一化结果的同时输出每样本每组的均值与倒标准差便于后续梯度计算复用。从实现角度看该算子是一个典型的「两趟扫描 分组并行」的归一化算子Tiling 侧按 group 粒度在多个 AI Core 间做负载均衡并依据 group 大小选择 small-group/general 两种调度模板Kernel 侧以双缓冲流水搬运数据通过 Kahan 补偿求和保证统计精度对非 float 类型统一提升到 float 计算后再回写。配套的 op_api 与 op_host 单元测试覆盖了正常路径、可选输入为空、空指针、非法数据类型、非法分组数与非法输出 shape 等场景为二次开发与调试提供了完整参照。如果需要在 NPU 上快速验证该算子可直接参照 test_aclnn_group_norm.cpp 中的用例组织输入输出张量并确保参数满足 README 中列出的产品支持与约束条件。【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表