CANN ops-nn BatchNorm 算子深度解析：训练/推理双模式实现与图模式调用实战指南

本文以 CANN（Compute Architecture for Neural Networks）神经网络算子库 ops-nn 中的 BatchNorm 算子为对象，系统讲解其在 NPU 上的批量归一化计算原理、训练与推理两种工作模式下的输入输出语义、全部可配置参数，并结合算子 IR 定义、Shape 推导、Tiling 计算与 Kernel 实现源码，给出通过算子 IR 构图（GEIR）方式在昇腾硬件上调用该算子的完整可运行样例。读完本文，你将掌握 BatchNorm 算子的参数语义、训练/推理差异、源码级实现路径，以及如何基于 算子 IR 定义 完成一次真实的图模式调用。

产品支持情况

根据 README 文档 的官方声明，BatchNorm 算子对以下产品系列均提供支持：

产品	是否支持
Ascend 950PR & 950DT 系列产品	√
Atlas A3 系列产品	√
Atlas A2 系列产品	√
Atlas 200I/500 A2 推理产品	√
Atlas 推理系列产品	√
Atlas 训练系列产品	√

从源码的 AICore 配置来看，batch_norm_def.cpp 中为 ascend950 与 ascend350 两代昇腾芯片注册了 AI Core 算子配置（OpAICoreConfig），而 op_kernel/batch_norm_apt.cpp 中对应的 arch35 目录存放了内核实现，这与上表所列产品系列（Ascend 950 系列、Atlas A3/A2 系列均基于相应芯片架构）相互印证。

功能说明与计算公式

算子功能

BatchNorm 算子对一个批次（batch）的数据执行批量归一化处理。归一化之后生成的数据其统计结果为 0 均值、1 标准差，即：y = (x - E(x)) / sqrt(Var(x) + ε) * γ + β。

具体而言，均值 E(x) 与方差 Var(x) 需要在算子内部计算得到；ε 是一个极小的浮点数，用于防止分母为 0 的情况；γ（scale）与 β（offset）是可学习的仿射参数。

计算公式

$$y = \frac{(x - E(x))}{\sqrt{Var(x) + ε}} * γ + β$$

其中各符号含义如下：

符号	含义
E(x)	输入 x 的均值（算子内部计算）
Var(x)	输入 x 的方差（算子内部计算）
ε	极小浮点数，防止除零，对应属性 epsilon，默认 1e-4f
γ	缩放权重，对应输入 scale
β	偏置，对应输入 offset

训练与推理两种模式

从 算子 IR 定义注释 可以确认，BatchNorm 通过属性 is_training 区分两种工作模式，两种模式的计算路径与输入要求有本质差异：

• 训练模式（is_training=true，默认）：算子内部计算当前批次的均值与方差用于归一化，并同步更新运行时均值/方差（moving mean/variance），为后续推理阶段保存统计量。
• 推理模式（is_training=false）：mean 与 variance 变为必选输入，直接使用训练阶段保存的统计量做归一化，不再计算批次统计量。

这一语义在 op_host/batch_norm_def.cpp 中得到印证：mean 与 variance 两个输入在 OpDef 中均被声明为 OPTIONAL（可选），其是否必填取决于运行模式。

参数说明

以下参数表完整继承自 README 文档，并结合 算子 IR 定义 与 OpDef 注册 补充了数据类型约束与默认值来源。

输入参数

参数名	输入/输出/属性	描述	数据类型	数据格式
x	输入	进行批量归一化的输入张量，对应公式中的 x。shape 维度至少是 2D。	FLOAT32、FLOAT16、BFLOAT16	NCHW/NHWC/NCDHW/NDHWC/ND
scale	输入	进行批量归一化的权重，对应公式中的 γ。一个 1D 张量，shape 与输入 x 的维度 C 相同；如果输入 x 是 ND 格式，shape 与输入 x 的第 2 维度相同。	FLOAT32	ND
offset	输入	进行批量归一化的偏置值，对应公式中的 β。一个 1D 张量，shape 与入参 scale 保持一致。	FLOAT32	ND
mean	可选输入	训练场景：训练期间动量更新前的均值，为可选输入；推理场景：推理期间使用的均值，为必选输入，对应公式中的 E(x)。一个 1D 张量，shape 与入参 scale 保持一致。	FLOAT32	ND
variance	可选输入	训练场景：训练期间动量更新前的方差，为可选输入；推理场景：推理期间使用的方差，为必选输入，对应公式中的 Var(x)。一个 1D 张量，shape 与入参 scale 保持一致。	FLOAT32	ND

属性（Attribute）

参数名	输入/输出/属性	描述	数据类型	数据格式
epsilon	可选属性	添加到方差中的小值以避免除以零，对应公式中的 ε。默认值为 1e-4f。	FLOAT32	-
data_format	可选属性	指定输入 x 的数据格式，支持 "NHWC"、"NCHW"、"NDHWC"、"NCDHW"。默认值为 "NHWC"。	STRING	-
is_training	可选属性	标记是否训练场景，true 表示训练场景，false 表示推理场景。默认值为 true。	BOOL	-
exponential_avg_factor	可选属性	训练期间用于更新动量平均值的因子。默认值为 1.0f。	FLOAT32	-

以上四个属性的默认值在 算子 IR 定义 中通过 REG_OP 的 ATTR 宏逐一登记：

.ATTR(epsilon, Float, 1e-4f)
.ATTR(data_format, String, "NHWC")
.ATTR(is_training, Bool, true)
.ATTR(exponential_avg_factor, Float, 1.0f)

同时在 OpDef 注册 中保持一致，其中 exponential_avg_factor 的默认常量定义在文件头部（DEFAULT_EXPONENTIAL_AVG_FACTOR = 1.0）。

输出参数

参数名	输入/输出/属性	描述	数据类型	数据格式
y	输出	表示批量归一化后的输出结果，对应公式中的 y。数据类型、数据格式、shape 与输入 x 保持一致。	FLOAT32、FLOAT16、BFLOAT16	NCHW/NHWC/NCDHW/NDHWC/ND
batch_mean	输出	训练模式：当前批次的均值；推理模式：等于输入 mean。一个 1D 张量，shape 与入参 scale 保持一致。	FLOAT32	ND
batch_variance	输出	训练模式：当前批次的方差；推理模式：等于输入 variance。一个 1D 张量，shape 与入参 scale 保持一致。	FLOAT32	ND
reserve_space_1	输出	为梯度计算预留。训练模式：等于输出 batch_mean；推理模式：等于输入 mean。一个 1D 张量，shape 与入参 scale 保持一致。	FLOAT32	ND
reserve_space_2	输出	为梯度计算预留。训练模式：保存 inv_var（1/sqrt(epsilon + variance)），用于反向梯度计算中重用；推理模式：等于输入 variance。一个 1D 张量，shape 与入参 scale 保持一致。	FLOAT32	ND
reserve_space_3	输出	一个只包含一个元素的一维张量，仅出于与 TensorFlow 兼容的目的而存在，数据无意义。shape 为 [1]。	FLOAT32	ND

类型约束细节：IR 定义中通过 DATATYPE 宏声明了两个类型模板——T1（作用于 x 与 y）支持 {DT_FLOAT, DT_FLOAT16, DT_BF16}，T2（作用于 scale、offset、mean、variance 及全部统计量输出）固定为 {DT_FLOAT}。也就是说，即使输入 x 是 FLOAT16/BFLOAT16，所有的通道级参数（scale/offset/mean/variance）以及 batch_mean、batch_variance、reserve_space_* 输出一律以 FLOAT32 参与计算，这与 batch_norm_infershape.cpp 中 InferDataType 的实现完全一致。

动量更新语义（源码级补充）

关于 exponential_avg_factor 与训练模式下动量统计量的更新，IR 定义注释 给出了明确的更新公式：

new_mean = (1 - factor) * old_mean + factor * batch_mean

• mean 输入在训练模式下作为动量均值的初始化起点：当提供该输入时用于初始化动量均值，未提供时从 0 开始；
• variance 输入在训练模式下作为动量方差的初始化起点：未提供时从 1 开始；
• 在 Kernel 实现 中可以看到，训练路径会根据 useRunningMeanVar（即是否传入 mean/variance）决定是否计算 oneSubMomentum = 1 - momentum，并在 CalculateRuningMeanVarVF 中完成动量均值和动量方差的更新；
• 训练模式下 batch_variance 的计算还带有 Bessel 校正（无偏估计）：besselCorrectionFactor = reduceNum / (reduceNum - 1)，见 batch_norm_full_reduce.h，这一点与 PyTorch 等框架在训练时使用无偏方差的口径一致。

约束说明

根据 README 文档，BatchNorm 算子无额外约束。

可以理解为：只要输入满足上表列出的数据类型、shape 维度和数据格式组合（例如 4D 张量使用 NCHW/NHWC，5D 张量使用 NCDHW/NDHWC，任意维度使用 ND），即可正常调用。

图模式调用方式

BatchNorm 算子支持图模式调用：通过算子 IR 构图方式调用 BatchNorm 算子，即直接基于 算子 IR 定义 中的 op::BatchNorm 构造计算图，再交由 GE（Graph Engine）完成构图、编译与执行。

调用方式	样例代码	说明
图模式调用	test_geir_batch_norm.cpp	通过算子 IR 构图方式调用 BatchNorm 算子。

样例整体流程

test_geir_batch_norm.cpp 演示了一次完整的 GEIR（Graph Engine Intermediate Representation）调用，其执行流程为：

1. 调用 ge::GEInitialize(global_options) 初始化 GE，设置 ge.exec.deviceId=0、ge.graphRunMode=1；
2. 创建 ge::Graph 计算图（图名 tc_ge_irrun_test）；
3. 在 CreateOppInGraph 中通过宏快速声明输入占位算子、常量算子并接入 op::BatchNorm("batchNorm_1")；
4. 为 BatchNorm 算子设置输入/输出 TensorDesc、属性并加入图；
5. 创建 ge::Session，通过 session->AddGraph(graph_id, graph, graph_options) 将图加入会话；
6. 通过 aclgrphDumpGraph 将图 dump 为 txt 便于检查；
7. 通过 session->RunGraph(graph_id, input, output) 执行图；
8. 将输入/输出数据写为 .bin 文件，并逐元素打印输出结果；
9. 通过 GEGetErrorMsgV2 / GEGetWarningMsgV2 获取运行期错误与告警信息，最后 GEFinalize 收尾。

构图核心代码解读

样例中的构图入口函数 CreateOppInGraph 定义了如下张量 shape 与属性（源码位于 test_geir_batch_norm.cpp）：

auto batchNorm_1 = op::BatchNorm("batchNorm_1");
std::vector<int64_t> xShape = {2, 3, 4, 5};        // 输入 x：4D，N=2, H=3, W=4, C=5
std::vector<int64_t> channelShape = {5};           // scale/offset/mean/variance：通道维度
std::vector<int64_t> oneShape = {1};               // reserve_space_3：单元素
ADD_INPUT(1, x, inDtype, xShape);
ADD_INPUT(2, scale, inDtype, channelShape);
ADD_INPUT(3, offset, inDtype, channelShape);
ADD_INPUT(4, mean, inDtype, channelShape);
ADD_INPUT(5, variance, inDtype, channelShape);
ADD_OUTPUT(1, y, inDtype, xShape);
ADD_OUTPUT(2, batch_mean, inDtype, channelShape);
ADD_OUTPUT(3, batch_variance, inDtype, channelShape);
ADD_OUTPUT(4, reserve_space_1, inDtype, channelShape);
ADD_OUTPUT(5, reserve_space_2, inDtype, channelShape);
ADD_OUTPUT(6, reserve_space_3, inDtype, oneShape);
ADD_INPUT_ATTR(epsilon, 1e-4f);
ADD_INPUT_ATTR(data_format, "NHWC");
ADD_INPUT_ATTR(is_training, true);
ADD_INPUT_ATTR(exponential_avg_factor, 1.0f);

其中：

• x 使用 4D shape {2, 3, 4, 5}，对应 data_format="NHWC" 时 N=2、H=3、W=4、C=5，通道数为 5；
• scale、offset、mean、variance 均为长度为 5（与通道数一致）的 1D 张量；
• 六个输出中 y 与 x 同 shape，batch_mean/batch_variance/reserve_space_1/reserve_space_2 与 scale 同 shape，reserve_space_3 为 [1]；
• 四个属性均使用默认值，即 epsilon=1e-4f、data_format="NHWC"、is_training=true、exponential_avg_factor=1.0f。

宏 ADD_INPUT 内部通过 op::Data 构造输入占位算子，并调用 GenOnesDataFloat32 将输入初始化为全 2 的 FLOAT32 数据；ADD_CONST_INPUT 则通过 op::Const 构造常量输入。这种宏封装的方式便于扩展更多输入规模，是理解 GEIR 构图的良好起点。

输出数据核对

样例在 RunGraph 成功后会遍历所有输出，将输出数据写入 ./tc_ge_irrun_test_0008_npu_output_*.bin，并逐元素打印。由于输入数据为常量（全 2），且训练模式下 batch_mean 即为当前批次的均值，读者可通过输出与手工计算（公式代入）互相印证算子的正确性。

底层实现原理（源码级扩展）

BatchNorm 算子在 ops-nn 仓库中的完整实现链路为：算子 IR 定义 → Shape/DataType 推导 → Tiling 计算 → Kernel 内核执行，下面依次展开。

1. 算子 IR 定义

op_graph/batch_norm_proto.h 使用 REG_OP(BatchNorm) 宏完成算子原语注册，声明了 5 个输入（其中 mean、variance 为 OPTIONAL_INPUT）、6 个输出与 4 个属性，并约束了 T1/T2 的数据类型集合。这部分直接决定算子对外暴露的接口语义，也是图模式调用（op::BatchNorm）的编译依据。

2. Shape 与 DataType 推导

op_host/batch_norm_infershape.cpp 实现 InferShape 与 InferDataType：

• y 的 shape 直接继承 x（*yShape = *xShape）；
• batch_mean、batch_variance、reserve_space_1、reserve_space_2 的 shape 直接继承 scale；
• reserve_space_3 固定为 gert::Shape({1})；
• 数据类型推导中 y 与 x 同类型，其余五个统计量输出固定为 DT_FLOAT。

该推导逻辑在 tests/ut/op_host/test_batch_norm_infershape.cpp 中有对应的单测覆盖：构造 x shape {2,3,4,5}、scale shape {5}，校验输出 0 为 [2, 3, 4, 5]、输出 1 为 [5]、输出 5 为 [1]，与推导实现一一对应。

3. Tiling 计算（多核切分）

op_host/arch35/batch_norm_tiling.cpp 注册了 Tiling4BatchNorm 与 TilingPrepare4BatchNorm：

• TilingPrepare 阶段从平台信息中获取 AIV 核数（coreNum）、UB 内存大小（ubSize）、UB 块大小（blockSize）与向量寄存器长度（vectorLength），用于后续的切分决策；
• 具体的切分实现分布在 op_host/arch35 目录下的多个文件中，覆盖 FullReduce、Welford 在线算法、BlockSplitR（按通道分块）、RA（多核归约） 等多种策略；
• 从 Kernel 入口 batch_norm_apt.cpp 可以看到，不同策略通过 TILINGKEY 进行标识（如 TILINGKEY_FULL_REDUCE=200000、TILINGKEY_WELFORD_REDUCE=300000、TILINGKEY_INFER=910000 等），Tiling 阶段选定的策略会在运行时通过 TILING_KEY_IS 分发到对应的 Kernel 实现类。

4. Kernel 内核实现

op_kernel/batch_norm_apt.cpp 定义了 AI Core 内核入口 batch_norm，按 TilingKey 分发到 11 种实现类，可归纳为两大族：

• 训练族：BatchNormFullReduce、BatchNormWelford、BatchNormBlockSplitR、BatchNormRAFullReduce、BatchNormRAWelford、BatchNormRARBlockSplitR 等，负责批次均值/方差计算（含 Welford 在线更新与 Bessel 校正）、动量统计量更新，以及归一化后输出 y；
• 推理族：BatchNormInfer、BatchNormInferLastChannel、BatchNormInferLastChannelContinuousA、BatchNormInferLastChannelSmallA、BatchNormInferSmallAB1 等，直接使用输入的 mean/variance 完成归一化。

以训练路径的 batch_norm_full_reduce.h 为例，其 ProcessUB 流程清晰展示了内核的计算管线：

1. CopyInX：将 x 按块从全局内存（GM）搬运到统一缓冲区（UB）；
2. CalculateMeanVarVF：在 UB 中计算当前批次的均值与方差（生成 batch_mean 与 batch_rstd，其中 rstd 即 1/sqrt(epsilon + variance)）；
3. CopyInGammaBeta / CopyInRunningMeanVar：加载 scale、offset 及动量均值/方差；
4. CalculateRuningMeanVarVF：完成动量均值和方差的更新；
5. CalculateNormalizeVF：执行 y = (x - batch_mean) * batch_rstd * gamma + beta；
6. 通过 EnQue/DeQue 双缓冲队列（DOUBLE_BUFFER=2）实现数据搬移与计算的重叠，最后 CopyOutY 写回全局内存。

此外，batch_norm_base.h（batch_norm_base.h）中针对 FLOAT32/FLOAT16/BFLOAT16 提供了统一的加载/存储模板：FLOAT16/BFLOAT16 数据先经 Cast 提升为 FLOAT32 参与运算，输出时再降回原精度，这也印证了"计算在 FLOAT32 精度下完成"的实现事实。

5. 二进制配置文件

op_host/config/ascend950/batch_norm_binary.json（以及 ascend350 版本）为每种"输入数据类型 × 数据格式"组合声明了独立的 kernel 二进制文件，例如 BatchNorm_bfloat16_NCDHW、BatchNorm_float16_NCHW、BatchNorm_float32_NHWC 等，共覆盖 bfloat16/float16/float32 三种输入类型与 NCDHW/NCHW/NHWC/NDHWC/ND 多种格式组合。配置中的 format_match_mode（FormatFixed/FormatDefault）进一步约束了输入输出格式的匹配规则，是算子编译期选核（kernel binary 选择）的直接依据。

总结

BatchNorm 是神经网络训练与推理中最基础的归一化算子之一。在 CANN ops-nn 算子库中，它通过 算子 IR 定义对外接口，以 is_training 属性区分训练（计算批次统计量 + 更新动量）与推理（使用输入统计量）两种模式；通过 Shape 推导、Tiling 多核切分 与 AI Core Kernel 构成完整的编译-执行链路；开发者可参照 test_geir_batch_norm.cpp 以图模式快速完成算子验证，并可通过 单测 与 ST 用例 深入核对推导与执行结果。