从零实现微型自动微分引擎：原理与实践

1. 从零实现微型自动微分引擎

周末突发奇想，决定动手实现一个微型PyTorch风格的自动微分引擎。这个项目我几年前曾经做过一次，但当时的代码早已不知所踪。这次我决定完整记录开发过程，从最基本的张量操作开始，逐步实现自动微分功能。

我们先从最基础的Tensor类开始。这个类本质上是对NumPy数组的简单封装，但会逐步加入自动微分所需的功能：

python复制import numpy as np

class Tensor:
    def __init__(self, data):
        self.data = data if isinstance(data, np.ndarray) else np.array(data)
        self._ctx = None  # 用于存储计算图的上下文信息

    def __add__(self, other):
        fn = Function(Add, self, other)
        result = Add.forward(self, other)
        result._ctx = fn
        return result

    def __mul__(self, other):
        fn = Function(Mul, self, other)
        result = Mul.forward(self, other)
        result._ctx = fn
        return result

    def __repr__(self):
        return f"tensor({self.data})"

注意：_ctx属性是自动微分的核心，它记录了张量是如何计算得到的，这在反向传播时会非常关键。

2. 计算图与自动微分原理

自动微分(Automatic Differentiation)是现代深度学习框架的核心技术。与符号微分和数值微分不同，自动微分通过记录计算过程来实现高效准确的微分计算。

2.1 前向计算与反向传播

在我们的实现中，每个Tensor对象都携带一个_ctx属性，它记录了：

• 产生该Tensor的操作(如Add或Mul)
• 参与计算的输入Tensor

这样，当我们从最终的输出Tensor开始，就可以沿着计算图回溯，逐步计算每个操作的梯度。

2.2 基本运算的微分规则

我们先实现加法和乘法的微分规则：

python复制class Function:
    def __init__(self, op, *args):
        self.op = op
        self.args = args

class Add:
    @staticmethod
    def forward(x, y):
        return Tensor(x.data + y.data)
    
    @staticmethod
    def backward(ctx, grad):
        x, y = ctx.args
        return Tensor([1]), Tensor([1])  # 加法的导数为1

class Mul:
    @staticmethod
    def forward(x, y):
        return Tensor(x.data * y.data)
    
    @staticmethod
    def backward(ctx, grad):
        x, y = ctx.args
        return Tensor(y.data), Tensor(x.data)  # 乘法的导数分别为y和x

提示：这里的backward方法接收两个参数 - ctx是计算上下文，grad是从上游传递来的梯度值。

3. 计算图可视化

理解计算图的结构对于调试自动微分系统非常重要。我们可以使用graphviz库来可视化计算图：

python复制import graphviz
from tinytorch import *

G = graphviz.Digraph(format='png')

def visit_nodes(G: graphviz.Digraph, node: Tensor):
    uid = str(id(node))
    G.node(uid, f"Tensor: {str(node.data)}")
    
    if node._ctx:
        ctx_uid = str(id(node._ctx))
        G.node(ctx_uid, f"Context: {str(node._ctx.op.__name__)}")
        G.edge(uid, ctx_uid)
        
        for child in node._ctx.args:
            G.edge(ctx_uid, str(id(child)))
            visit_nodes(G, child)

if __name__ == "__main__":
    x = Tensor([8])
    y = Tensor([5])
    z = x + y
    visit_nodes(G, z)
    G.render(directory="vis", view=True)

这段代码会生成一个PNG图像，清晰地展示从输入Tensor到输出Tensor的计算路径。

4. 反向传播实现

现在我们已经准备好了实现反向传播的所有组件。反向传播的核心思想是链式法则 - 从输出开始，沿着计算图反向传播梯度。

python复制def backward(tensor, grad=None):
    if grad is None:
        grad = Tensor(np.ones_like(tensor.data))  # 默认梯度为1
    
    if tensor._ctx is None:  # 叶子节点
        return
    
    ctx = tensor._ctx
    op = ctx.op
    
    # 计算当前操作的梯度
    grads = op.backward(ctx, grad)
    
    # 递归反向传播
    for arg, g in zip(ctx.args, grads):
        backward(arg, g)

让我们测试一下这个实现：

python复制if __name__ == "__main__":
    x = Tensor([2.0])
    y = Tensor([3.0])
    z = x * y
    
    backward(z)
    
    print(f"dz/dx: {x.grad}")  # 应该输出3.0
    print(f"dz/dy: {y.grad}")  # 应该输出2.0

5. 常见问题与调试技巧

在实现自动微分引擎时，经常会遇到一些典型问题：

5.1 梯度计算错误

症状：反向传播得到的梯度值与预期不符。
排查步骤：

1. 检查计算图是否正确构建
2. 验证每个操作的backward实现是否正确
3. 确保梯度传递没有中断

5.2 内存泄漏

症状：长时间运行后内存占用持续增长。
解决方案：

• 及时清理不再需要的计算图节点
• 实现梯度清零机制

5.3 数值不稳定

症状：梯度计算中出现NaN或inf。
解决方案：

• 添加梯度裁剪
• 检查数值范围是否合理

6. 扩展功能与性能优化

基础版本完成后，可以考虑添加以下功能提升实用性：

6.1 支持更多运算

python复制class Sub:
    @staticmethod
    def forward(x, y):
        return Tensor(x.data - y.data)
    
    @staticmethod
    def backward(ctx, grad):
        return Tensor([1]), Tensor([-1])

class Div:
    @staticmethod
    def forward(x, y):
        return Tensor(x.data / y.data)
    
    @staticmethod
    def backward(ctx, grad):
        x, y = ctx.args
        return Tensor(1/y.data), Tensor(-x.data/y.data**2)

6.2 矩阵运算支持

python复制class MatMul:
    @staticmethod
    def forward(x, y):
        return Tensor(np.dot(x.data, y.data))
    
    @staticmethod
    def backward(ctx, grad):
        x, y = ctx.args
        return Tensor(np.dot(grad.data, y.data.T)), Tensor(np.dot(x.data.T, grad.data))

6.3 GPU加速

可以通过CuPy替换NumPy来实现GPU加速：

python复制try:
    import cupy as cp
    use_gpu = True
except ImportError:
    use_gpu = False
    import numpy as cp

class Tensor:
    def __init__(self, data):
        self.data = cp.array(data)
        self._ctx = None

7. 实际应用示例

让我们用这个微型框架实现一个简单的线性回归：

python复制# 生成随机数据
np.random.seed(42)
X = np.random.rand(100, 1)
y = 3 * X + 2 + 0.1 * np.random.randn(100, 1)

# 初始化参数
w = Tensor(np.random.randn(1, 1))
b = Tensor(np.zeros(1))

# 训练循环
learning_rate = 0.1
for epoch in range(100):
    # 前向传播
    pred = X @ w + b
    loss = ((pred - y) ** 2).mean()
    
    # 反向传播
    backward(loss)
    
    # 参数更新
    w.data -= learning_rate * w.grad
    b.data -= learning_rate * b.grad
    
    # 梯度清零
    w._ctx = None
    b._ctx = None
    
    if epoch % 10 == 0:
        print(f"Epoch {epoch}, Loss: {loss.data}")

这个简单的例子展示了如何使用我们的自动微分引擎训练机器学习模型。虽然功能有限，但核心原理与主流框架相同。

实现过程中最关键的收获是理解了计算图的构建和梯度传播机制。自动微分看似复杂，但拆解后其实是由许多简单的部分组成的。下一步可以考虑添加更多优化器、损失函数和网络层，逐步完善这个微型框架。