手搓自动微分
技术背景
自动微分技术,在各大深度学习框架里面得到了广泛的应用。但是其实究其原理,就是一个简单的链式法则。要实现一个自动微分框架是非常容易的事情,难的是高阶的自动微分和端到端的自动微分。这篇文章主要介绍一阶自动微分的基础Python实现,以及一些简单的测试案例。
链式法则
求导的链式法则,这个在高数里面大家就都学过了,形式比较简单:
或者可以写成这种形式:
自动微分框架的使用
我们先用一些现成的自动微分框架,如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]
这里面的函数定义为:
其导数解析形式为:
也可以用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等自动微分框架计算出来的结果进行比对,可以看到结果都是一致的。
总结概要
不同于符号微分、手动微分和差分法,自动微分方法有着使用简单、计算精度较高、性能较好等优势,因此在各大深度学习框架中得到了广泛的应用。虽然每个框架所使用的自动微分的原理不尽相同,但大致都是基于链式法则计算结合图计算的一些优化。如果是自己动手来手搓一个自动微分框架的话,大致就只能实现一下一阶的链式法则的自动微分。
版权声明
本文首发链接为:https://www.cnblogs.com/dechinphy/p/auto-grad.html
作者ID:DechinPhy
更多原著文章:https://www.cnblogs.com/dechinphy/
请博主喝咖啡:https://www.cnblogs.com/dechinphy/gallery/image/379634.html