手搓自动微分

DechinPhy

发布于 2024-04-19 08:23:30

1150

发布于 2024-04-19 08:23:30

技术背景

自动微分技术，在各大深度学习框架里面得到了广泛的应用。但是其实究其原理，就是一个简单的链式法则。要实现一个自动微分框架是非常容易的事情，难的是高阶的自动微分和端到端的自动微分。这篇文章主要介绍一阶自动微分的基础Python实现，以及一些简单的测试案例。

链式法则

求导的链式法则，这个在高数里面大家就都学过了，形式比较简单：

f(g(x))'=f'[g(x)]\cdot g'(x)

或者可以写成这种形式：

\frac{df}{dx}=\frac{df}{dg}\cdot\frac{dg}{dx}

自动微分框架的使用

我们先用一些现成的自动微分框架，如MindSpore，来演示一下自动微分的基本用法：

import numpy as np
from mindspore import grad, Tensor
from mindspore import numpy as msnp
# 定义一个自变量x
x = Tensor(np.array([1., 2., 3.], np.float32))
# 定义一个复合函数
f = lambda x: msnp.sin(msnp.cos(x))
# 函数求导
gf = grad(f)
# 计算自动微分结果
print (gf(x))
# [-0.7216062  -0.831692   -0.07743199]

这里面的函数定义为：

f(x) = \sin(\cos(x))

其导数解析形式为：

f'(x)=-\cos(\cos(x))\sin(x)

也可以用MindSpore做一个简单的验证：

print (-msnp.cos(msnp.cos(x))*msnp.sin(x))
# [-0.7216062  -0.831692   -0.07743199]

可以看到结果是一致的。

手搓自动微分

自己实现自动微分，其实就是把每一个操作函数的导数函数定义好，例如我们可以定义某一个操作的求导函数为__grad__()，而求值函数在python中有一个内置的__call__()函数。例如我们可以基于numpy的函数自定义一个正弦函数的类：

import numpy as np
class SIN:
    def __call__(self, x):
        """计算正弦值"""
        return np.sin(x)
    def __grad__(self, x):
        """计算正弦函数的导数值"""
        return np.cos(x)

然后配套一个grad自动微分函数：

def grad(obj):
    """直接调用输入操作的自动微分函数"""
    return obj.__grad__

甚至可以实现一个value_and_grad函数，同时计算值和导数：

class ValueAndGrad:
    def __init__(self, obj):
        """初始化输入对象的求值函数和求导函数"""
        self.obj1 = obj
        self.obj2 = obj.__grad__
    def __call__(self, x):
        """用元组的形式将值和导数的计算结果返回"""
        return (self.obj1(x), self.obj2(x))
def value_and_grad(obj):
    """初始化求值求导对象"""
    return ValueAndGrad(obj)

需要注意的是，因为大多数的场景下都会涉及到复合函数的计算，这也是自动微分技术的核心之一，因此我们自己实现的自动微分框架要能够接收一些外来的操作，然后在内部递归的计算。对应的带有自动微分的类格式变为：

class SIN:
    def __init__(self, obj=None):
        """给定一个其他的函数"""
        self.obj = obj
    def __call__(self, x):
        """没有复合函数时直接返回结果，有复合函数就递归计算"""
        return np.sin(x) if self.obj is None else np.sin(self.obj(x))
    def __grad__(self, x):
        """没有复合函数时直接返回导数结果，有复合函数就按照链式法则递归计算"""
        return COS()(x) if self.obj is None else COS()(self.obj(x))*self.obj.__grad__(x)

最终形成的自动微分实现案例为：

import numpy as np
import mindspore as ms
from mindspore import Tensor
from mindspore import grad as msgrad
from mindspore import numpy as msnp

class SIN:
    """自定义正弦类"""
    def __init__(self, obj=None):
        self.obj = obj
    def __call__(self, x):
        return np.sin(x) if self.obj is None else np.sin(self.obj(x))
    def __grad__(self, x):
        return COS()(x) if self.obj is None else COS()(self.obj(x))*self.obj.__grad__(x)
    
class COS:
    """自定义余弦类"""
    def __init__(self, obj=None):
        self.obj = obj
    def __call__(self, x):
        return np.cos(x) if self.obj is None else np.cos(self.obj(x))
    def __grad__(self, x):
        return -SIN()(x) if self.obj is None else -SIN()(self.obj(x))*self.obj.__grad__(x)

class ValueAndGrad:
    """自定义求值求导类"""
    def __init__(self, obj):
        self.obj1 = obj
        self.obj2 = obj.__grad__
    def __call__(self, x):
        return (self.obj1(x), self.obj2(x))

def grad(obj):
    """自定义求导函数"""
    return obj.__grad__

def value_and_grad(obj):
    """自定义求值求导函数"""
    return ValueAndGrad(obj)

# 定义自变量
x = np.array([0., 1., 2., 3.,], np.float32)
# 单体函数验证
assert np.allclose(SIN()(x), np.sin(x))
# 单体函数求导验证
assert np.allclose(grad(SIN())(x), np.cos(x))
v, g = value_and_grad(SIN())(x)
# 单体函数求值求导验证
assert np.allclose(v, np.sin(x))
assert np.allclose(g, np.cos(x))
# 双复合函数验证
assert np.allclose(SIN(SIN())(x), np.sin(np.sin(x)))
assert np.allclose(SIN(COS())(x), np.sin(np.cos(x)))
assert np.allclose(COS(SIN())(x), np.cos(np.sin(x)))
assert np.allclose(COS(COS())(x), np.cos(np.cos(x)))
# 三复合函数验证
assert np.allclose(SIN(COS(SIN()))(x), np.sin(np.cos(np.sin(x))))
# 双复合函数求导验证
assert np.allclose(grad(SIN(SIN()))(x), np.cos(x)*np.cos(np.sin(x)))
tensor_x = Tensor(x, ms.float32)
ms_func1 = lambda x: msnp.sin(msnp.cos(x))
assert np.allclose(grad(SIN(COS()))(x), msgrad(ms_func1)(tensor_x).asnumpy())
ms_func2 = lambda x: msnp.cos(msnp.sin(x))
assert np.allclose(grad(COS(SIN()))(x), msgrad(ms_func2)(tensor_x).asnumpy())
ms_func3 = lambda x: msnp.cos(msnp.sin(msnp.cos(x)))
# 三复合函数求导验证
assert np.allclose(grad(COS(SIN(COS())))(x), msgrad(ms_func3)(tensor_x).asnumpy())

这里面除了可以跟手推的微分解析形式的计算结果进行比对之外，还可以跟MindSpore等自动微分框架计算出来的结果进行比对，可以看到结果都是一致的。

总结概要

不同于符号微分、手动微分和差分法，自动微分方法有着使用简单、计算精度较高、性能较好等优势，因此在各大深度学习框架中得到了广泛的应用。虽然每个框架所使用的自动微分的原理不尽相同，但大致都是基于链式法则计算结合图计算的一些优化。如果是自己动手来手搓一个自动微分框架的话，大致就只能实现一下一阶的链式法则的自动微分。