自动微分autograd实践要点

本文主要参考 反向传播和神经网络训练 · 大神Andrej Karpathy 的“神经网络从Zero到Hero 系列”之一，提炼一些精要，将反向传播的细节和要点展现出来

定义Value

第一步首先要定义Value，算子中需包含：data（Value 的数值）， grad（Value 的梯度），_backward（反向传播函数，初始化为 None），_prev（需要依赖于它的Value，用于后面构建反向传播链）：

class Value:
  
  def __init__(self, data, _children=(), _op='', label=''):
    self.data = data
    self.grad = 0.0
    self._backward = lambda: None
    self._prev = set(_children)
    self._op = _op
    self.label = label

  def __repr__(self):
    return f"Value(data={self.data})"
  
  def __add__(self, other):
    out = Value(self.data + other.data, (self, other), '+')
    
    def _backward():
      self.grad += 1.0 * out.grad
      other.grad += 1.0 * out.grad
    out._backward = _backward
    
    return out

  def __mul__(self, other):
    out = Value(self.data * other.data, (self, other), '*')
    
    def _backward():
      self.grad += other.data * out.grad
      other.grad += self.data * out.grad
    out._backward = _backward
      
    return out
  
  def tanh(self):
    x = self.data
    t = (math.exp(2*x) - 1)/(math.exp(2*x) + 1)
    out = Value(t, (self, ), 'tanh')
    
    def _backward():
      self.grad += (1 - t**2) * out.grad
    out._backward = _backward
    
    return out
  
  def backward(self):
    
    topo = []
    visited = set()
    def build_topo(v):
      if v not in visited:
        visited.add(v)
        for child in v._prev:
          build_topo(child)
        topo.append(v)
    build_topo(self)
    
    self.grad = 1.0
    for node in reversed(topo):
      node._backward()

上述代码的核心在于：

1. 每个算子的_backward 函数需要依次按算子进行手动定义
2. 一个Value的backward函数，是从当前Value开始，先将依赖于这个Value的所有Value按依赖顺序串起来，然后再从当前Value开始依次运行_backward()。这样只需要对一条链上的最后一个Value运行backward函数，就可以将这个链上的所有节点的grad全更新一次，即完成一次反向传播

例如：

a = Value(2.0, label='a')
b = Value(-3.0, label='b')
c = Value(10.0, label='c')
e = a*b; e.label = 'e'
d = e + c; d.label = 'd'
f = Value(-2.0, label='f')
L = d * f; L.label = 'L'

上述计算链条的图示为：

---

手动定义每个 operator 的 _backward() 函数

加法
不同算子的梯度值不一样，例如：对于加法：out = Value(self.data + other.data, (self, other), '+') ，out 对 self. Data 求导，倒数值为1，因此其_backward() 函数定义为：

def _backward():
      self.grad += 1.0 * out.grad
      other.grad += 1.0 * out.grad

这里是 += 而不是 = 的原因是，有的时候某个node的Value，在前向传播时可能影响了不止一个Value，例如：

那么这里 a.grad 即要计算从 d 处来的反向传播，也要考虑从 e 处来的反向传播，因此是 +=

乘法
同理，对于乘法 out = Value(self.data * other.data, (self, other), '*')，out 对 self. Data 求导，倒数值为 other. Data，因此：

def _backward():
      self.grad += other.data * out.grad
      other.grad += self.data * out.grad

激活函数
tanh 激活函数为：
$$tanh(x)=\frac{e^x-e^{-x}}{e^x+e^{-x}}=\frac{e^{2x}-1}{e^{2x}+1}$$
其倒数为 $$tan{h}^{\prime }(x)=1-(\frac{e^{2x}-1}{e^{2x}+1}{)}^2=1-(tanh(x){)}^2$$
因此对应的公式为

def tanh(self):
    x = self.data
    t = (math.exp(2*x) - 1)/(math.exp(2*x) + 1)
    out = Value(t, (self, ), 'tanh')
    
    def _backward():
      self.grad += (1 - t**2) * out.grad

---

构建反向传播计算链

以上图为例：

反向传播就是从 L 开始从右往左依次调用各个Node的 _backward()，因此链条构建的方式类似于树的遍历，从根节点开始往逐渐添加叶节点：

def backward(self):
    
    topo = []
    visited = set()
    def build_topo(v):
      if v not in visited:
        visited.add(v)
        for child in v._prev:
          build_topo(child)
        topo.append(v)
    build_topo(self)
    
    self.grad = 1.0
    for node in reversed(topo):
      node._backward()

这里根节点对自身进行求导，倒数值都为1，所以需要设置self.grad = 1.0。最后只需运行一次 L.backward() 就可以把所有Node的梯度全更新一遍，以下是运行一次L.backward()后的结果：

---

Reference：

1. 反向传播和神经网络训练 · 大神Andrej Karpathy 的“神经网络从Zero到Hero 系列”之一