mindspore.ops.rms_norm

查看源文件
mindspore.ops.rms_norm(x, gamma, epsilon=1e-6)[源代码]

RmsNorm(Root Mean Square Layer Normalization)即均方根标准化。与LayerNorm相比,其保留了缩放不变性,而舍弃了平移不变性。 其公式如下:

\[y=\frac{x_i}{\sqrt{\frac{1}{n}\sum_{i=1}^{n}{ x_i^2}+\varepsilon}}\gamma_i\]

警告

这是一个实验性API,后续可能修改或删除。该API目前只支持在Atlas A2训练系列产品上使用。

参数:
  • x (Tensor) - RmsNorm的输入, 支持的数据类型为: float16、float32、bfloat16。

  • gamma (Tensor) - 可训练参数(\(\gamma\)),支持的数据类型: float16、float32、bfloat16。

  • epsilon (float, 可选) - 一个取值范围为(0, 1]的浮点值,用于避免除零。默认值: 1e-6

返回:
  • Tensor,归一化后的结果,shape和数据类型与 x 相同。

  • Tensor,类型为float,表示输入数据标准差的倒数,用于反向梯度计算。

异常:
  • TypeError - x 的数据类型不是float16、float32、bfloat16中的一种。

  • TypeError - gamma 的数据类型不是float16、float32、bfloat16中的一种。

  • TypeError - xgamma 的数据类型不一致。

  • ValueError - epsilon 不是一个0到1之间的float值。

  • ValueError - gamma 的秩大于 x 的秩。

支持平台:

Ascend

样例:

>>> import mindspore
>>> import numpy as np
>>> from mindspore import Tensor, ops
>>> x = Tensor(np.array([[1, 2, 3], [1, 2, 3]]), mindspore.float32)
>>> gamma = Tensor(np.ones([3]), mindspore.float32)
>>> y, rstd = ops.rms_norm(x, gamma)
>>> print(y)
[[0.46290997  0.92581993  1.3887299]
 [0.46290997  0.92581993  1.3887299]]
>>> print(rstd)
[[0.46290997]
 [0.46290997]]