MindSpore梯度进阶操作

时间：2024-05-16　|　来源：　|　阅读：73

话题： S in

这篇文章主要介绍了mindspore深度学习框架中基于InsertGradientOf算子的进阶梯度操作。InsertGradientOf算子的功能跟此前介绍过的bprop功能有些类似，也是自定义梯度，但bprop更倾向于计算梯度，而InsertGradientOf算子更倾向于修改梯度，这里介绍了一

技术背景

在MindSpore深度学习框架中，我们可以使用 mindspore.grad 对函数式编程的函数直接计算自动微分，也可以使用 mindspore.ops.GradOperation 求解Cell类的梯度dout。本文所介绍的 mindspore.ops.InsertGradientOf 是一个对dout进一步进行处理的算子，类似于在Cell类中自定义一个bprop函数，不改变前向传播输出的结果，只改变反向传播的结果。

测试场景

我们使用一个简单的函数 \(f(x,y)=xy^2,\frac{\partial f}{\partial x}=y^2\) 来测试一下MindSpore中的自动微分，以及InsertGradientOf算子对梯度的操作。

import numpy as np
from mindspore import Tensor, ops, grad
# 定义Clip函数的上下界
a = Tensor(np.array([1.1]).astype(np.float32))
b = Tensor(np.array([0.1]).astype(np.float32))
# Clip反向传播dx结果
def clip_gradient(dx):
    ret = dx
    if ret > a:
        ret = a
    if ret < b:
        ret = b
    return ret
# 生成一个计算图的节点
clip = ops.InsertGradientOf(clip_gradient)
# 主函数
func = lambda x, y: x * y ** 2
# 带Clip的主函数
def clip_func(x, y):
    x = clip(x)
    c = func(x, y)
    return c
# 带Clip主函数的前向传播
def f(x, y):
    return clip_func(x, y)
# 带Clip主函数的反向传播
def fd(x, y):
    return grad(clip_func)(x, y)
# 给定x，y的数值
x = Tensor(np.array([-2]).astype(np.float32))
y = Tensor(np.array([2]).astype(np.float32))
# 分别计算Clip前后的主函数的前向传播与反向传播
print("forward: ", func(x, y))
print("backward: ", grad(func)(x, y))
print("clip forward: ", f(x, y))
print("clip backward:", fd(x, y))

输出结果为：

forward:  [-8.] # x*y**2
backward:  [4.] # y**2
clip forward:  [-8.] # x*y**2
clip backward: [1.1] # clip(y**2)

需要注意的是，虽然我们最终clip的时候操作的是 \(\frac{\partial f}{\partial x}=y^2\) ，但是在函数实现时，clip函数应该施加在 \(x\) 上面，而不是 \(y\) 上面，这表示对 \(x\) 的反向传播进行操作。

InsertGradientOf成员函数

bprop是MindSpore框架中Cell类的一个关于计算反向传播的函数，可以用于计算和处理梯度值。但是有一个比较偏的问题是，bprop的函数输入与construct函数的输入要求要一致，如果参数数量对不上，就会报错。关于这一点，其实torch里面处理的方案会更直观一些，可以参考这篇博客中的两个Issue。而MindSpore中要实现类似的功能，就需要依赖于这个InsertGradientOf算子。先看一个使用bprop处理Clip梯度的示例：

import numpy as np
from mindspore import Tensor, grad, nn
# 定义Clip参数
a = Tensor(np.array([1.1]).astype(np.float32))
b = Tensor(np.array([0.1]).astype(np.float32))
# Clip函数
def clip_gradient(dx):
    ret = dx
    if ret > a:
        ret = a
    if ret < b:
        ret = b
    return ret
# 需要被Clip的Cell类
class Net(nn.Cell):
    # 使用bprop处理梯度
    def bprop(self, x, y, out, dout):
        return clip_gradient(y**2)
    def construct(self, x, y):
        return x * y ** 2

x = Tensor(np.array([-2]).astype(np.float32))
y = Tensor(np.array([2]).astype(np.float32))
net = Net()
print (net(x, y))
print (grad(net)(x, y))

输出结果为：

[-8.]
[1.1]

这里还是比较容易理解的，我们手动推导了一个 \(\frac{\partial f}{\partial x}=y^2\) ，那么就可以把 \(y\) 参数传给bprop函数，然后计算 \(y^2\) ，最后再计算clip。但是这个方案要求传入到bprop函数的参数是完整的，如果参数匹配不上就会报错：

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

a = Tensor(np.array([1.1]).astype(np.float32))
b = Tensor(np.array([0.1]).astype(np.float32))

def clip_gradient(dx):
    ret = dx
    if ret > a:
        ret = a
    if ret < b:
        ret = b
    return ret

clip = ops.InsertGradientOf(clip_gradient)

class Net(nn.Cell):
    def bprop(self, y, out, dout): 
        return clip_gradient(y ** 2)
    def construct(self, x, y):
        return x * y ** 2

x = Tensor(np.array([-2]).astype(np.float32))
y = Tensor(np.array([2]).astype(np.float32))

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

这里给bprop函数传入的参数跟construct函数是对不齐的，那么计算梯度时就会出现这样的报错：

[-8.]
Traceback (most recent call last):
  File "test_insert_gradient.py", line 83, in 
    print (grad(net)(x, y))
  File "/home/dechin/anaconda3/envs/mindspore-latest/lib/python3.7/site-packages/mindspore/ops/composite/base.py", line 622, in after_grad
    return grad_(fn_, weights, grad_position)(*args, **kwargs)
  File "/home/dechin/anaconda3/envs/mindspore-latest/lib/python3.7/site-packages/mindspore/common/api.py", line 131, in wrapper
    results = fn(*arg, **kwargs)
  File "/home/dechin/anaconda3/envs/mindspore-latest/lib/python3.7/site-packages/mindspore/ops/composite/base.py", line 601, in after_grad
    res = self._pynative_forward_run(fn, grad_, weights, args, kwargs)
  File "/home/dechin/anaconda3/envs/mindspore-latest/lib/python3.7/site-packages/mindspore/ops/composite/base.py", line 658, in _pynative_forward_run
    outputs = fn(*args, **new_kwargs)
  File "/home/dechin/anaconda3/envs/mindspore-latest/lib/python3.7/site-packages/mindspore/nn/cell.py", line 693, in __call__
    raise err
  File "/home/dechin/anaconda3/envs/mindspore-latest/lib/python3.7/site-packages/mindspore/nn/cell.py", line 690, in __call__
    _pynative_executor.end_graph(self, output, *args, **kwargs)
  File "/home/dechin/anaconda3/envs/mindspore-latest/lib/python3.7/site-packages/mindspore/common/api.py", line 1264, in end_graph
    self._executor.end_graph(obj, output, *args, *(kwargs.values()))
TypeError: Size of bprop func inputs[1] is not equal to the size of cell inputs[2]

----------------------------------------------------
- C++ Call Stack: (For framework developers)
----------------------------------------------------
mindspore/ccsrc/pipeline/pynative/grad/grad.cc:837 GetCustomBpropPrim

但是我们知道， \(\frac{\partial f}{\partial x}=g(y)\) 是一个只跟 \(y\) 有关的函数，其实不用传入 \(x\) 参数也应该要可以计算其梯度值。

接下来考虑，如果在Cell类外定义一个InsertGradientOf算子构建的函数，那么也可以在Cell类里面使用：

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

a = Tensor(np.array([1.1]).astype(np.float32))
b = Tensor(np.array([0.1]).astype(np.float32))

def clip_gradient(dx):
    ret = dx
    if ret > a:
        ret = a
    if ret < b:
        ret = b
    return ret

clip = ops.InsertGradientOf(clip_gradient)

class Net(nn.Cell):
    def construct(self, x, y):
        return clip(x) * y ** 2

x = Tensor(np.array([-2]).astype(np.float32))
y = Tensor(np.array([2]).astype(np.float32))

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

输出结果为：

[-8.]
[1.1]

这个计算结果是对的，不过我们需要的是这个clip函数最好也能够调用到类本身的一些属性和成员变量，而InsertGradientOf算子也支持对成员函数进行处理：

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

a = Tensor(np.array([1.1]).astype(np.float32))
b = Tensor(np.array([0.1]).astype(np.float32))

class Net(nn.Cell):
    def __init__(self):
        super().__init__()
        self.clip = ops.InsertGradientOf(self.back)
    # 把Clip定义成一个成员函数
    def back(self, y): 
        ret = y
        if ret > a:
            ret = a
        if ret < b:
            ret = b
        return ret
    def construct(self, x, y):
        return self.clip(x) * y ** 2

x = Tensor(np.array([-2]).astype(np.float32))
y = Tensor(np.array([0.8]).astype(np.float32))

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

这里输出的结果为：

[-1.2800001]
[0.64000005]

因为 \(y^2=0.64\) ，未触发边界Clip的条件，因此这里正常输出 \(\frac{\partial f}{\partial x}=y^2\) ，如果稍微调整下输入的 \(y\) ，触发了边界条件，那么梯度就会被Clip：

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

a = Tensor(np.array([1.1]).astype(np.float32))
b = Tensor(np.array([0.1]).astype(np.float32))

class Net(nn.Cell):
    def __init__(self):
        super().__init__()
        self.clip = ops.InsertGradientOf(self.back)
    def back(self, y): 
        ret = y
        if ret > a:
            ret = a
        if ret < b:
            ret = b
        return ret
    def construct(self, x, y):
        return self.clip(x) * y ** 2

x = Tensor(np.array([-2]).astype(np.float32))
y = Tensor(np.array([2]).astype(np.float32))

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

计算结果为：

[-8.]
[1.1]

当然了，如果我们直接返回一个跟 \(x\) 、 \(y\) 都无关的参数作为梯度也是可以的：

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

a = Tensor(np.array([1.1]).astype(np.float32))
b = Tensor(np.array([0.1]).astype(np.float32))

class Net(nn.Cell):
    def __init__(self):
        super().__init__()
        self.clip = ops.InsertGradientOf(self.back)
    def back(self, dx): 
        return 100.
    def construct(self, x, y):
        return self.clip(x) * y ** 2

x = Tensor(np.array([-2]).astype(np.float32))
y = Tensor(np.array([2]).astype(np.float32))

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

输出结果为：

[-8.]
100.0

如果要再传一些偏置参数到 \(x\) 的梯度中，例如令 \(g=\frac{\partial f}{\partial x}+z\) ，而这个参数 \(z\) 一般都是通过construct函数直接传进Cell类的。此时可用的思路是，把这些额外的变量存到类的属性里面，通过读取成员变量再加载到梯度操作函数中：

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

a = Tensor(np.array([1.1]).astype(np.float32))
b = Tensor(np.array([0.1]).astype(np.float32))

class Net(nn.Cell):
    def __init__(self):
        super().__init__()
        self.clip = ops.InsertGradientOf(self.back)
        self.z = 0.
    def back(self, dx): 
        ret = dx
        if ret > a:
            ret = a
        if ret < b:
            ret = b
        return ret + self.z
    def construct(self, x, y, z=0.):
        self.z = z
        return self.clip(x) * y ** 2

x = Tensor(np.array([-2]).astype(np.float32))
y = Tensor(np.array([2]).astype(np.float32))

net = Net()
print (net(x, y, z=-1))
print (grad(net)(x, y, z=-1))

输出结果为：

[-8.]
[0.10000002]

这就实现了给梯度修饰函数传参的功能。

优先级问题

凡是有冲突的操作，就必然有一个优先级的顺序。bprop函数是用本地的方法去计算一个梯度值，而InsertGradientOf算子是对某一个变量的梯度值进行处理。因此当这两个函数同时被用于处理一个梯度值时，就需要看看谁的优先级更高：

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

a = Tensor(np.array([1.1]).astype(np.float32))
b = Tensor(np.array([0.1]).astype(np.float32))

class Net(nn.Cell):
    def __init__(self):
        super().__init__()
        self.clip = ops.InsertGradientOf(self.back)
    def back(self, y): 
        ret = y
        if ret > a:
            ret = a
        if ret < b:
            ret = b
        return ret
    def bprop(self, x, y, out, dout):
        return 100.
    def construct(self, x, y):
        return self.clip(x) * y ** 2

x = Tensor(np.array([-2]).astype(np.float32))
y = Tensor(np.array([2]).astype(np.float32))

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

在这个案例中，clip函数还是对梯度做一个截断，而bprop函数则是直接返回一个梯度值。那么最终执行的输出结果为：