自动微分autograd实践要点

news2024/11/28 4:42:34

定义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()

上述代码的核心在于：

每个算子的_backward 函数需要依次按算子进行手动定义
一个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 激活函数为：
$\frac{e^x -e^{-x}}{e^x + e^{-x}} = \frac{e^{2x} -1}{e^{2x} + 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()