【MindSpore入门教程】 02 自动微分

深度学习等现代AI算法，自动微分技术（Automatic Differentiation，AD）是其中的关键技术。

自动微分是一种介于数值微分与符号微分之间的一种求导方法。自动微分的核心思想是将计算机程序中的运算操作分解为一个有限的基本操作合集，且合集中基本操作的求导规则均为已知的。在完成每一个基本操作的求导后，使用链式求导法则将结果组合得到整体程序的求导结果。
链式求导法则：
$$(f\circ g{)}^{{}^{\prime }}(x)=f^{{}^{\prime }}(g(x))g^{{}^{\prime }}(x)\text{\hspace{1em}(1)}$$

自动微分实现的方式可以分为前项自动微分和反向自动微分，MindSpore的反向算子GradOperation实现的是反向自动微分。所以这里主要介绍一下反向自动微分。反向自动微分的求导方向与原函数的求值方向相反，微分结果需要依赖原函数的运行结果。

反向自动微分

我们将被求导的原函数，泛化为具有N输入与M输出的函数F：
$$(Y_1,Y_2,...,Y_M)=F(X_1,X_2,...,X_N)\text{\hspace{1em}(2)}$$

函数 F 的导数本身为一个雅可比矩阵(Jacobian matrix)。
$$\begin{array}{r}\left[\begin{array}{ccc}\frac{\partial Y_1}{\partial X_1}& ...& \frac{\partial Y_1}{\partial X_N}\\ ...& ...& ...\\ \frac{\partial Y_M}{\partial X_1}& ...& \frac{\partial Y_M}{\partial X_N}\end{array}\right]\end{array}\text{\hspace{1em}(3)}$$

在反向自动微分当中，我们是从输出开始向输入的方向计算的，因此每一次计算我们可以求得某一输出对输入的导数，即雅可比矩阵中的一行。
$$\begin{array}{r}\left[\begin{array}{ccc}\frac{\partial Y_i}{\partial X_1}& ...& \frac{\partial Y_i}{\partial X_N}\end{array}\right]\end{array}\text{\hspace{1em}(4)}$$

为了求取该列的值$\frac{\partial Y_i}{\partial X_1}$， 自动微分将程序分解为一系列求导规则已知的基本操作，这些基本操作也可以被泛化表达为具有n输入和m输出的函数 f：
$$(y_1,y_2,...,y_m)=f(x_1,x_2,...,x_n)\text{\hspace{1em}(5)}$$

由于我们的已知基础函数 f 求导规则，即其雅可比矩阵是已知的。于是我们可以对 f 计算向量雅可比积(Vjp, Vector-jacobian-product)，并应用链式求导法则获得导数结果。
$$\begin{array}{r}\left[\begin{array}{ccc}\frac{\partial Y_i}{\partial x_1}& ...& \frac{\partial Y_i}{\partial x_N}\end{array}\right]=\left[\begin{array}{ccc}\frac{\partial Y_i}{\partial y_1}& ...& \frac{\partial Y_i}{\partial y_m}\end{array}\right]\left[\begin{array}{ccc}\frac{\partial y_1}{\partial x_1}& ...& \frac{\partial y_1}{\partial x_n}\\ ...& ...& ...\\ \frac{\partial y_m}{\partial x_1}& ...& \frac{\partial y_m}{\partial x_n}\end{array}\right]\end{array}\text{\hspace{1em}(6)}$$
其中，$\frac{\partial Y_i}{\partial y_j}$，$\frac{\partial y_j}{\partial x_k}$ 就是基本函数的求导。

GradOperation实现

单算子的反向

为简单说明一个算子的微分，举例正弦函数的微分，已知：
$$\begin{gathered} y=f(x)=sin(x) \\ {y}^{\prime }=f^{\prime }(x)=cos(x)\text{\hspace{1em}(7)} \end{gathered}$$

如下，即是在MindSpore中实现正弦函数的微分。
其中，ops.Cos()是余弦函数，x是输入变量，out是sinx(x)计算的结果，即正向输出结果。dout表示的是链式法则中已经计算过得微分。很多基础算子的反向求导，是通过装饰器@bprop_getters.register(xxx)进行关联的。

import mindspore.ops as ops
from mindspore.ops._grad.grad_base import bprop_getters

@bprop_getters.register(ops.Sin)
def get_bprop_sin(self):
    """Grad definition for `Sin` operation."""
    cos = ops.Cos()
    def bprop(x, out, dout):
        dx = dout * cos(x)
        return (dx,)
    return bprop

有了基础算子的反向导数，即可以根据链式法则，组合符合算子的反向。

复合网络的反向

对于一个简单的网络
$$f(x)=cos(sin(x))\text{\hspace{1em}(8)}$$
其对输入x进行求导
$$f^{\prime }(x)=-sin(sin(x))\ast cos(x)\text{\hspace{1em}(9)}$$
首先可以使用MindSpore定义正向的网络f(x) :

from mindspore import nn, ops
from mindspore import Tensor
import numpy as np

class Net(nn.Cell):
    def __init__(self):
        super(Net, self).__init__()
        self.sin = ops.Sin()
        self.cos = ops.Cos()

    def construct(self, x):
        a = self.sin(x)
        out = self.cos(a)
        return out

def forward():
    f = Net()
    x = np.array([1.0, 2.0, 5.0], np.float32)
    y1 = np.cos(np.sin(x))
    x = Tensor(x)
    y2 = f(x)

    print("y1 = ", y1)
    print("y2 = ", y2)

if __name__ == "__main__":
    forward()

# output
#y1 =  [0.6663667  0.6143003  0.57440084]
#y2 =  [0.66636676 0.6143003  0.57440084]

前向计算结果一致。
MindSpore提供了GradOperation高阶函数，用于构造输入函数的反向微分：

from mindspore import nn, ops
from mindspore import Tensor
import numpy as np

class GradNet(nn.Cell):
	def __init__(self, net):
        super(GradNet, self).__init__()
        self.net = net
        self.grad_op = ops.GradOperation()
    def construct(self, x):
        gradient_function = self.grad_op(self.net)
        return gradient_function(x)

def backward():
    f = Net()
    gradf = GradNet(f)
    x = np.array([1.0, 2.0, 5.0], np.float32)
    y1 = -np.sin(np.sin(x)) * np.cos(x)
    x = Tensor(x)
    y2 = gradf(x)

    print("y1 = ", y1)
    print("y2 = ", y2)

if __name__ == "__main__":
    backward()

#output 
#y1 =  [-0.40286243  0.32836995  0.23219854]
#y2 =  [-0.4028624   0.32836998  0.23219854]

可以发现，首先使用GradOperation算子，计算出网络的微分函数gradient_function = self.grad_op(self.net)，然后，使用微分函数进行反向计算，结果与预期结果一致。

GradOperation使用

GradOperation有三个输入参数get_all, get_by_list, sens_param，包含多种使用方式。其中，

• get_all=True， 表示对网络的所有参数(required=True)进行求导。

• get_by_list=True, 表示对网络中指定的参数进行求导，此时计算网络的微分函数，需要传入参数列表： gradient_function = self.grad_op(self.net, self.params).

• sens_param=True, 表示配置（关于输出的梯度的）灵敏度，可以手动控制输出梯度值的大小。反向计算时需要向梯度函数传递一个灵敏度值grad_wrt_output。这个输入值必须与正向网络的输出具有相同的形状和类型。使用方式如 gradient_function(x, y, grad_wrt_output) 。

下面看几个例子。

1. 生成一个梯度函数，该函数返回第一个输入的梯度

from mindspore import Parameter
# f(x,y,z) = (x*z)*y
class Net(nn.Cell):
    def __init__(self):
        super(Net, self).__init__()
        self.matmul = ops.MatMul()
        self.z = Parameter(Tensor(np.array([1.0], np.float32)), name='z')
    def construct(self, x, y):
        x = x * self.z
        out = self.matmul(x, y)
        return out
# df/dx
class GradNetWrtX(nn.Cell):
    def __init__(self, net):
        super(GradNetWrtX, self).__init__()
        self.net = net
        self.grad_op = ops.GradOperation()
    def construct(self, x, y):
        gradient_function = self.grad_op(self.net)
        return gradient_function(x, y)

x = Tensor([[0.5, 0.6, 0.4], [1.2, 1.3, 1.1]], dtype=mstype.float32)
y = Tensor([[0.01, 0.3, 1.1], [0.1, 0.2, 1.3], [2.1, 1.2, 3.3]], dtype=mstype.float32)

grad = GradNetWrtX(Net())(x, y)
print(grad)

# out = (x*z)*y 关于x导数如下
ynp= y.asnumpy().T
a = np.ones((2,3))
out = np.matmul(a, ynp)
print(out)

#output
#[[1.4100001 1.5999999 6.6      ]
# [1.4100001 1.5999999 6.6      ]]
#[[1.41000004 1.59999996 6.5999999 ]
# [1.41000004 1.59999996 6.5999999 ]]

关于矩阵的求导，可以参考矩阵求导公式的数学推导。

2. 生成一个梯度函数，该函数返回所有输入的梯度

# df/dx, df/dy
class GradNetWrtXY(nn.Cell):
    def __init__(self, net):
        super(GradNetWrtXY, self).__init__()
        self.net = net
        self.grad_op = ops.GradOperation(get_all=True)
    def construct(self, x, y):
        gradient_function = self.grad_op(self.net)
        return gradient_function(x, y)

x = Tensor([[0.5, 0.6, 0.4], [1.2, 1.3, 1.1]], dtype=mstype.float32)
y = Tensor([[0.01, 0.3, 1.1], [0.1, 0.2, 1.3], [2.1, 1.2, 3.3]], dtype=mstype.float32)
output = GradNetWrtXY(Net())(x, y)
print(output)

# out = (x*z)*y 关于x导数同上
# out = (x*z)*y 关于y导数如下
xnp = x.asnumpy().T
a = np.ones((2,3))
out = np.matmul(xnp, a)
print(out)

#output
#(Tensor(shape=[2, 3], dtype=Float32, value=
#[[1.41000009e+000, 1.59999990e+000, 6.59999990e+000],
# [1.41000009e+000, 1.59999990e+000, 6.59999990e+000]]), 
# Tensor(shape=[3, 3], dtype=Float32, value=
#[[1.70000005e+000, 1.70000005e+000, 1.70000005e+000],
# [1.89999998e+000, 1.89999998e+000, 1.89999998e+000],
# [1.50000000e+000, 1.50000000e+000, 1.50000000e+000]]))
#[[1.70000005 1.70000005 1.70000005]
# [1.89999998 1.89999998 1.89999998]
# [1.50000003 1.50000003 1.50000003]]

关于矩阵的求导，可以参考矩阵求导公式的数学推导。

3. 生成一个梯度函数，该函数返回给定参数的梯度

1. 构造一个带有 get_by_list=True 参数的高阶函数： grad_op = GradOperation(get_by_list=True)。
2. 创建一个 ParameterTuple 和 net 作为GradOperation 的高阶函数的参数输入，ParameterTuple作为参数过滤器决定返回哪个梯度：params = ParameterTuple(net.trainingable_params()) 。
3. 使用高阶函数，将net和 params 作为参数输入 GradOperation 的高阶函数，得到梯度函数： gradient_function = grad_op(net, params) 。

from mindspore import ParameterTuple
# df/dz
class GradNetWithWrtParams(nn.Cell):
    def __init__(self, net):
        super(GradNetWithWrtParams, self).__init__()
        self.net = net
        self.params = ParameterTuple(net.trainable_params())
        self.grad_op = ops.GradOperation(get_by_list=True)
    def construct(self, x, y):
        gradient_function = self.grad_op(self.net, self.params)
        return gradient_function(x, y)

x = Tensor([[0.5, 0.6, 0.4], [1.2, 1.3, 1.1]], dtype=mstype.float32)
y = Tensor([[0.01, 0.3, 1.1], [0.1, 0.2, 1.3], [2.1, 1.2, 3.3]], dtype=mstype.float32)
output = GradNetWithWrtParams(Net())(x, y)
print(output)

# out = (x*z)*y 关于z导数如下
#output
#(Tensor(shape=[1], dtype=Float32, value= [1.53369999e+001]),)

out = (x*z)*y 矩阵关于z的导数，不知道怎么计算的。有人知道的请多指教。

4. 生成一个梯度函数，该函数返回的梯度乘以一个灵敏度

1. 构建一个带有 get_all=True, sens_param=True 参数的 GradOperation 高阶函数：grad_op = GradOperation(get_all=True, sens_param=True) 。
2. 当 sens_param=True ，定义 grad_wrt_output （关于输出的梯度）：grad_wrt_output = Tensor([[0.5, 0.5, 0.5], [0.5, 0.5, 0.5]], dtype=mstype.float32) 。
3. 用 net 作为参数输入 ，得到梯度函数：gradient_function = grad_op(net) 。
4. 用 net的输入和 sens_param 作为参数调用梯度函数，得到关于所有输入的梯度：gradient_function(x, y, grad_wrt_output) 。

import mindspore
# sens * df/dx, sens * df/dy, sens * df/dz
class GradNetWrtXYWithSensParam(nn.Cell):
    def __init__(self, net):
        super(GradNetWrtXYWithSensParam, self).__init__()
        self.net = net
        self.grad_op = ops.GradOperation(get_all=True, get_by_list=True, sens_param=True)
        self.params = ParameterTuple(net.trainable_params())
        self.grad_wrt_output = Tensor([[0.1, 0.1, 0.1], [0.1, 0.1, 0.1]], dtype=mindspore.float32)
    def construct(self, x, y):
        gradient_function = self.grad_op(self.net, self.params)
        return gradient_function(x, y, self.grad_wrt_output)
  
x = Tensor([[0.5, 0.6, 0.4], [1.2, 1.3, 1.1]], dtype=mindspore.float32)
y = Tensor([[0.01, 0.3, 1.1], [0.1, 0.2, 1.3], [2.1, 1.2, 3.3]], dtype=mindspore.float32)

output = GradNetWrtXYWithSensParam(Net())(x, y)
print(output)

# out = (x*z)*y 关于y导数如下
xnp = x.asnumpy().T
a = np.ones((2,3))
out = np.matmul(xnp, a)
print(out)

#output
"""
((Tensor(shape=[2, 3], dtype=Float32, value=
[[1.41000003e-001, 1.59999996e-001, 6.59999967e-001],
 [1.41000003e-001, 1.59999996e-001, 6.59999967e-001]]), 
 Tensor(shape=[3, 3], dtype=Float32, value=
[[1.70000002e-001, 1.70000002e-001, 1.70000002e-001],
 [1.89999998e-001, 1.89999998e-001, 1.89999998e-001],
 [1.50000006e-001, 1.50000006e-001, 1.50000006e-001]])), 
 (Tensor(shape=[1], dtype=Float32, value= [1.53369999e+000]),))
[[1.70000005 1.70000005 1.70000005]
 [1.89999998 1.89999998 1.89999998]
 [1.50000003 1.50000003 1.50000003]]
"""

与上面反向计算结果对比，网络反向求解的关于输入x,y和参数z的导数，都乘上了0.1。 这就是灵敏度参数的作用，自定义调节反向梯度的大小。

特别地，在设置了sens_param=True时，一定要记住在gradient_function中输入一个灵敏度Tensor，且其shape大小要求与正向计算的结果保证一致。否则，可能有如下报错信息：

# 错误代码
return gradient_function(x, y)  # 缺少self.grad_wrt_output

# 报错信息
Traceback (most recent call last):
  File "gradoperation.py", line 75, in <module>
    output = GradNetWrtXYWithSensParam(Net())(x, y)
  File "mindspore\lib\site-packages\mindspore\nn\cell.py", line 578, in __call__
    out = self.compile_and_run(*args)
  File "mindspore\lib\site-packages\mindspore\nn\cell.py", line 965, in compile_and_run
    self.compile(*inputs)
  File "mindspore\lib\site-packages\mindspore\nn\cell.py", line 938, in compile
    jit_config_dict=self._jit_config_dict)
  File "mindspore\lib\site-packages\mindspore\common\api.py", line 1134, in compile
    result = self._graph_executor.compile(obj, args_list, phase, self._use_vm_mode())
TypeError: The parameters number of the function is 3, but the number of provided arguments is 2.
FunctionGraph ID : construct.25
NodeInfo: In file gradoperation.py(13)
    def construct(self, x, y):
    
# 错误代码
self.grad_wrt_output = Tensor([[0.1, 0.1], [0.1, 0.1]], dtype=mindspore.float32) # shape与正向输出不一致

# 错误信息
[ERROR] ANALYZER(169608,1,?):2022-8-7 19:43:42 [mindspore\ccsrc\pipeline\jit\static_analysis\async_eval_result.cc:66] HandleException] Exception happened, check the information as below.

The function call stack (See file 'rank_0\om/analyze_fail.dat' for more details. Get instructions about `analyze_fail.dat` at https://www.mindspore.cn/search?inputValue=analyze_fail.dat):
# 0 In file gradoperation.py(70)
        return gradient_function(x, y, self.grad_wrt_output)
               ^
# 1 In file mindspore\lib\site-packages\mindspore\ops\_grad\grad_math_ops.py(250)
        if ta:
# 2 In file mindspore\lib\site-packages\mindspore\ops\_grad\grad_math_ops.py(253)
            dx = mul1(dout, w)
                 ^

Traceback (most recent call last):
  File "gradoperation.py", line 75, in <module>
    output = GradNetWrtXYWithSensParam(Net())(x, y)
  File "mindspore\lib\site-packages\mindspore\nn\cell.py", line 578, in __call__
    out = self.compile_and_run(*args)
  File "mindspore\lib\site-packages\mindspore\nn\cell.py", line 965, in compile_and_run
    self.compile(*inputs)
  File "mindspore\lib\site-packages\mindspore\nn\cell.py", line 938, in compile
    jit_config_dict=self._jit_config_dict)
  File "mindspore\lib\site-packages\mindspore\common\api.py", line 1134, in compile
    result = self._graph_executor.compile(obj, args_list, phase, self._use_vm_mode())
  File "mindspore\lib\site-packages\mindspore\ops\primitive.py", line 467, in __check__
    fn(*(x[track] for x in args))
  File "mindspore\lib\site-packages\mindspore\ops\operations\math_ops.py", line 1436, in check_shape
    raise ValueError(f"For '{cls_name}', the input dimensions must be equal, but got 'x1_col': {x1_col} "
ValueError: For 'MatMul', the input dimensions must be equal, but got 'x1_col': 2 and 'x2_row': 3. And 'x' shape [2, 2](transpose_a=False), 'y' shape [3, 3](transpose_b=True).

至此，我们熟悉了MindSpore自动微分的设计原理，对使用方式也有了初步的了解。

参考资料

[1] mindspore.ops.GradOperation

[2] MindSpore 函数式微分编程

[3] Baydin, A.G. et al., 2018. Automatic differentiation in machine learning: A survey. arXiv.org. Available at: https://arxiv.org/abs/1502.05767