CANN/ops-nn PowerSign优化器原地更新算子

📅 2026/7/15 18:12:31
CANN/ops-nn PowerSign优化器原地更新算子
InplaceApplyPowerSign【免费下载链接】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 训练系列产品×功能说明算子功能执行PowerSign优化器的单步参数更新。PowerSign将更新机制解耦为方向和幅度两个维度利用梯度与一阶矩的符号一致性来自适应调整更新步长对权重var和一阶矩m进行原地更新inplace语义。对标TensorFlowtf.raw_ops.ApplyPowerSign/tf.raw_ops.ResourceApplyPowerSign接口。计算公式$$ \begin{aligned} m_{t} \beta \cdot m_{t-1} (1 - \beta) \cdot grad \ sign_gm \text{sign}(m_{t}) \cdot \text{sign}(grad) \ var_{t} var_{t-1} - lr \cdot \exp(logbase \cdot sign_decay \cdot sign_gm) \cdot grad \end{aligned} $$其中sign(x)为符号函数x 0 返回 1x 0 返回 -1x 0 返回 0exp()为自然指数函数。参数说明参数名输入/输出/属性描述数据类型数据格式var输入待更新的权重参数张量对应公式中的var。与图输出端口var共享GM地址inplace更新。BFLOAT16、FLOAT16、FLOATNDm输入一阶矩估计张量对应公式中的m。shape/dtype必须与var一致与图输出端口m共享GM地址inplace更新。BFLOAT16、FLOAT16、FLOATNDlr输入学习率对应公式中的lr。shape必须为 [1] 的scalar Tensor。BFLOAT16、FLOAT16、FLOATNDlogbase输入对数基底对应公式中的logbase。shape必须为 [1] 的scalar Tensor。BFLOAT16、FLOAT16、FLOATNDsign_decay输入符号衰减因子对应公式中的sign_decay。shape必须为 [1] 的scalar Tensor。BFLOAT16、FLOAT16、FLOATNDbeta输入一阶矩衰减率对应公式中的β取值范围 [0, 1)。shape必须为 [1] 的scalar Tensor。BFLOAT16、FLOAT16、FLOATNDgrad输入当前梯度张量对应公式中的grad。shape/dtype必须与var一致。BFLOAT16、FLOAT16、FLOATNDuse_locking属性是否在更新时加锁。默认false。当前实现不强制互斥锁仅作语义占位。BOOL-var输出更新后的权重张量与输入var共享Device内存inplace。shape/dtype与输入var完全相同。BFLOAT16、FLOAT16、FLOATNDm输出更新后的一阶矩估计张量与输入m共享Device内存inplace。shape/dtype与输入m完全相同。BFLOAT16、FLOAT16、FLOATND约束说明var、m、grad三个张量的shape必须一致。lr、logbase、sign_decay、beta必须为shape[1]的scalar Tensor。所有输入输出的DataType必须一致。FP16/BF16输入内部提升至FP32计算后cast回原类型。调用说明调用方式调用样例说明图模式test_geir_inplace_apply_power_sign通过算子IR构图方式调用InplaceApplyPowerSign算子。【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考