ApplyKerasMomentum

产品支持情况

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

功能说明

• 算子功能：执行Keras Momentum优化器的单步参数更新。根据动量系数momentum、学习率lr和梯度grad更新动量累积量accum，并按标准模式或Nesterov模式原地更新权重参数var（inplace 语义）。对标TensorFlow中tf.raw_ops.ResourceApplyKerasMomentum接口的计算语义。

• 计算公式：

  $$\begin{array}{rl}accu{m}_{new}& =momentum\cdot accum-lr\cdot grad\end{array}$$

  • 标准模式（use_nesterov = 0）：

  $$va{r}_{new}=var+accu{m}_{new}$$

  • Nesterov模式（use_nesterov = 1）：

  $$va{r}_{new}=var+momentum\cdot accu{m}_{new}-lr\cdot grad$$

  其中momentum为动量系数，lr为学习率，grad为当前梯度，accum为动量累积量，var为待更新的权重参数。

参数说明

参数名	输入/输出/属性	描述	数据类型	数据格式
var	输入 / 输出 (inplace)	待更新的权重参数，对应公式中的var。Kernel内inplace更新，GE IR输出var_out与输入var共享Device内存。	FLOAT16、FLOAT、BFLOAT16	ND
accum	输入 (inplace 更新)	动量累积量，对应公式中的accum。shape/dtype必须与var一致；Kernel内显式写回输入GM地址。	FLOAT16、FLOAT、BFLOAT16	ND
lr	输入	学习率，对应公式中的lr。shape={1}的1元素scalar Tensor，dtype为FLOAT。	FLOAT	ND
grad	输入	当前梯度Tensor，对应公式中的grad。shape/dtype必须与var一致。	FLOAT16、FLOAT、BFLOAT16	ND
momentum	输入	动量系数，对应公式中的momentum。shape={1}的1元素scalar Tensor，dtype为FLOAT。	FLOAT	ND
var	输出	更新后的var Tensor，与输入var共享Device内存（inplace）。	FLOAT16、FLOAT、BFLOAT16	ND
use_locking	属性	是否对更新操作加锁保护。Bool类型属性，默认值为false。	Bool	—
use_nesterov	属性	是否使用Nesterov动量，对应公式中的use_nesterov。Bool类型属性，默认值为false（false=标准模式，true=Nesterov模式）。	Bool	—

注：accum 为 in-place 更新（与 TensorFlow ResourceApplyKerasMomentum 一致），不作为显式输出返回，计算结果直接写回输入地址。

约束说明

• var、accum、grad必须具有相同的数据类型和形状。
• lr、momentum为标量Tensor（shape={1}），数据类型为FLOAT。
• use_locking、use_nesterov为Bool类型属性，默认值为false。

调用说明

调用方式	调用样例	说明
图模式	test_geir_apply_keras_momentum	通过 算子IR 构图方式调用 ApplyKerasMomentum 算子。