CANN/ops-nn ReLU6梯度算子

张开发
2026/6/10 21:33:49 15 分钟阅读

分享文章

CANN/ops-nn ReLU6梯度算子
Relu6Grad【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn产品支持情况产品是否支持Ascend 950PR/Ascend 950DT√Atlas A3 训练系列产品/Atlas A3 推理系列产品×Atlas A2 训练系列产品/Atlas A2 推理系列产品×Atlas 200I/500 A2 推理产品×Atlas 推理系列产品×Atlas 训练系列产品×本算子目前仅在 Ascend 950 (dav-3510 / Atlas 350 加速卡) 架构上实现。功能说明算子功能ReLU6 激活函数的反向梯度。对前向输入落在严格开区间(0, 6)的位置透传上游梯度dy落在区间外含端点x 0与x 6的位置输出 0。计算公式$$ dx_i \begin{cases} dy_i, 0 x_i 6 \ 0, \text{otherwise} \end{cases} $$严格开区间。x 0与x 6均落入零臂。x NaN由于 IEEE-754 比较语义返回 0x ±inf同样落入零臂。当 mask 为 true 时dy按值 透传dy NaN/Inf不会被屏蔽。第三方对标等价于 TensorFlowRelu6Grad与 PyTorchhardtanh_backward(min0, max6)在边界x ∈ {0, 6}取 0 的约定。参数说明参数名输入/输出描述数据类型数据格式gradients输入上游梯度dy。FLOAT16, FLOAT, BFLOAT16NDfeatures输入前向输入x非前向输出与gradients须可广播。FLOAT16, FLOAT, BFLOAT16NDbackprops输出反向梯度dx。shape 为gradients、features广播后的统一形状。FLOAT16, FLOAT, BFLOAT16ND约束说明gradients、features、backprops必须为同一种 dtype不支持 mix-dtype。不支持 INT32 等整型数据。支持任意 NumPy 广播形态含标量[1]、单维 broadcast、跨 rank broadcast。支持动态 shape 与动态 rank。features应当为 ReLU6 的前向输入x而不是前向输出。该约定与canndev/ops/built-in/tbe/impl/dynamic/relu6_grad.py一致。实现方案层文件说明计算图原型op_graph/relu6_grad_proto.hREG_OP(Relu6Grad)二输入一输出算子定义op_host/relu6_grad_def.cppOpDef::AddConfig(ascend950, ...)InferShapeop_host/relu6_grad_infershape.cpp复用Ops::Base::InferShape4BroadcastTilingop_host/arch35/relu6_grad_tiling_arch35.{h,cpp}按 dtype 分支调用Ops::Base::BroadcastBaseTilingOpDagDAGop_kernel/arch35/relu6_grad_dag.hfp32 通路原生计算fp16/bf16 通路提升 fp32 中间精度Structop_kernel/arch35/relu6_grad_struct.hBRC_TEMP_SCH_MODE_KEY_DECL/SELKernel 入口op_kernel/relu6_grad_apt.cppKERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_AIV_ONLY)BroadcastSchschMode, OpDagfp32 通路In0(dy) ─CopyInBrc─ ────────────────────────────── | In1(x) ─CopyInBrc─ | ─Compare(GT, x, 0) ─ mask_gt0 | ─Compare(LT, x, 6) ─ mask_lt6 | mask And(mask_gt0, mask_lt6) | Const(0,fp32) ─Duplicate─ Zero ─────────────── ┴ Select(mask, dy, Zero) ─ CopyOut ─ Out0Vec::Compareu8, fp32, GT/LT输出位掩码Vec::Andu8合成(0 x 6)区间 maskVec::Selectu8, fp32, TENSOR_TENSOR按mask ? dy : 0逐元素 选择避免分支硬件原生 vsel 实现。fp16 / bf16 通路In0/In1 ─CopyInBrc─ Cast(-fp32) ─ ─ Compare ── And ─ Select ── Cast(-T,RINT) ─ CopyOut ─ Out0 Const(0,fp32) ── Zero ─ Const(6,fp32) ── Six ──把 fp16 / bf16 整体提升到 fp32 做 cmp/and/sel末端用CAST_MODE_RINTround-to-nearest-even回退到原 dtype。原因与 DSLrelu6_grad.py在 fp16 vcmpsel 不可用时 fallback 到 fp32 的行为一致上下界 0.0 和 6.0 在 fp16/bf16 中均能精确表达但 cmp 走 fp32 路径更稳健 且 dy 是 NaN/Inf 时 fp32 中间不会被半精度范围截断。调用说明调用方式样例代码说明图模式test_geir_relu6_grad通过算子IR构图方式调用 Relu6Grad 算子覆盖 fp32/fp16/bf16 基础用例 边界 x0/x6、xNaN/±Inf、广播等关键特殊值用例。【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考

更多文章