nanograd：从零开始构建的微型深度学习框架

nanograd 是一个极简的深度学习框架，完全从零开始用 Python 实现了反向模式自动微分和神经网络层。本文深入解析其设计原理、核心机制以及教育价值。

现代深度学习框架如 PyTorch 和 TensorFlow 已经将复杂的数学运算和自动微分机制封装得极为完善，开发者往往只需几行代码就能训练复杂的神经网络。然而，这种高度的抽象也带来了问题——许多从业者虽然能够调用 API 完成工作，却对底层的自动微分原理、计算图构建以及反向传播机制缺乏直观理解。

nanograd 项目正是为了解决这一痛点而诞生的。它是一个从零开始用纯 Python 实现的微型深度学习框架，旨在通过亲手实现核心机制，帮助开发者真正理解现代机器学习框架背后的工作原理。

nanograd 是一个极简化但功能完整的深度学习框架，其核心目标是展示自动微分和神经网络层的最小可行实现。项目不依赖任何外部深度学习库，完全基于 NumPy 风格的张量操作和 Python 的闭包机制构建计算图。

尽管代码量精简，nanograd 却实现了现代深度学习框架的核心组件：反向模式自动微分（即反向传播）、张量操作、神经网络层、激活函数、损失函数以及优化器。这种"麻雀虽小，五脏俱全"的设计使其成为学习深度学习底层原理的理想教材。

自动微分是深度学习框架的基石。nanograd 实现了反向模式自动微分，这是训练神经网络时最高效的微分方式。其核心思想是通过构建计算图，在前向传播时记录所有操作，然后在反向传播时利用链式法则从输出节点向输入节点传播梯度。

nanograd 使用 Python 闭包来构建计算图。每个张量操作不仅返回计算结果，还会记录该操作的前置节点和梯度计算函数。当调用 backward() 方法时，框架通过拓扑排序确保按正确顺序计算梯度，并处理多分支情况下的梯度累加。

框架支持丰富的张量运算，包括：

• 逐元素运算：加法、减法、乘法、除法、幂运算
• 矩阵运算：矩阵乘法、转置、重塑
• 索引与切片：支持张量的切片和索引操作
• 归约操作：求和、均值、最大值等聚合运算
• 广播机制：自动处理不同形状张量之间的运算

这些操作为构建复杂神经网络提供了基础构件。

nanograd 实现了多种核心神经网络层：

线性层（Linear）：全连接层，实现 y = xW + b 的线性变换，是学习特征表示的基础。

卷积层（Conv1D/Conv2D）：一维卷积采用滑动窗口实现，二维卷积则使用 im2col 技术将卷积运算转换为矩阵乘法，这是工业界常用的优化策略。

激活函数：支持 ReLU、Tanh、Sigmoid 等常用激活函数，为网络引入非线性能力。

框架实现了两种优化算法：

随机梯度下降（SGD）：最基础的优化方法，沿负梯度方向更新参数，简单但有效。

BFGS 算法：一种拟牛顿二阶优化方法，利用近似海森矩阵信息指导参数更新。相比 SGD，BFGS 通常能在更少的迭代次数内收敛，但每次迭代的计算成本更高。项目中还实现了基于中心差分的数值海森计算，用于验证和调试。

使用 nanograd 训练一个简单的 XOR 分类器非常直观：

from tensors import Tensor
from layers import Linear, Sequential, Tanh
from losses import MSE
from optim import SGD

# 构建模型
model = Sequential(
    Linear(2, 4),
    Tanh(),
    Linear(4, 1)
)

criterion = MSE()
optimizer = SGD(model.parameters(), lr=0.1)

# 训练数据
x = Tensor([[0., 0.], [0., 1.], [1., 0.], [1., 1.]])
y = Tensor([[0.], [1.], [1.], [0.]])

# 训练循环
for _ in range(2000):
    pred = model(x)
    loss = criterion(y, pred)
    model.zero_grad()
    loss.backward()
    optimizer.step()

更令人印象深刻的是，使用 BFGS 优化器解决同样的问题仅需约 40 步迭代即可收敛，展示了二阶优化方法的高效性。