调用 PyTorch 训练模型并不难，真正困难的是回答这些问题：一次张量运算如何进入计算图？梯度为什么能沿图反向传播？线性层、激活函数和注意力机制最终又如何落到这些基础操作上？
TinyTorch 的思路不是再包装一层现成框架，而是把框架本身变成学习对象：从张量出发，逐步构建自动微分、神经网络模块、优化器，最终走向 Transformer。对于想理解机器学习系统内部机制的开发者，这种“自己造一个能工作的教学版本”的方式，比单纯记忆 API 更直接。
成熟框架隐藏了哪些关键步骤
在 PyTorch 中，一段训练代码通常只有几行：
loss = criterion(model(x), y)
loss.backward()
optimizer.step()
这几行背后至少包含四套机制：
- 张量存储与运算：管理形状、数据类型以及逐元素计算。
- 计算图：记录结果由哪些输入和操作产生。
- 自动微分：按照链式法则，从损失反向计算每个参数的梯度。
- 参数更新：优化器读取梯度并修改权重。
成熟框架必须处理 GPU、混合精度、分布式训练、算子融合和内存优化，因此源码很难直接作为第一份教材。教学框架的价值在于主动缩小问题：先在 CPU 上实现正确语义，再观察这些语义如何组合成完整模型。
一个可运行的微型自动微分引擎
下面可以这样实践：使用 NumPy 保存数据，为每次运算记录父节点和局部反向函数。它不是 TinyTorch 的源码复刻，而是一个可独立运行的最小练习，用来理解张量、计算图和反向传播如何协作。
先准备环境：
python -m venv .venv
source .venv/bin/activate
python -m pip install numpy
Windows PowerShell 可将激活命令替换为：
.venv\Scripts\Activate.ps1
将下面代码保存为 tiny_autograd.py：
import numpy as np
def sum_to_shape(grad, shape):
# 撤销广播：把梯度压缩回原始操作数的形状。
while grad.ndim > len(shape):
grad = grad.sum(axis=0)
for axis, size in enumerate(shape):
if size == 1 and grad.shape[axis] != 1:
grad = grad.sum(axis=axis, keepdims=True)
return grad
class Tensor:
def __init__(self, data, parents=(), backward=None):
self.data = np.asarray(data, dtype=np.float64)
self.grad = np.zeros_like(self.data)
self.parents = tuple(parents)
self._backward = backward or (lambda: None)
@staticmethod
def ensure(value):
return value if isinstance(value, Tensor) else Tensor(value)
def __add__(self, other):
other = Tensor.ensure(other)
out = Tensor(self.data + other.data, (self, other))
def backward():
self.grad += sum_to_shape(out.grad, self.data.shape)
other.grad += sum_to_shape(out.grad, other.data.shape)
out._backward = backward
return out
__radd__ = __add__
def __neg__(self):
return self * -1.0
def __sub__(self, other):
return self + (-Tensor.ensure(other))
def __mul__(self, other):
other = Tensor.ensure(other)
out = Tensor(self.data * other.data, (self, other))
def backward():
self.grad += sum_to_shape(out.grad * other.data, self.data.shape)
other.grad += sum_to_shape(out.grad * self.data, other.data.shape)
out._backward = backward
return out
__rmul__ = __mul__
def mean(self):
out = Tensor(self.data.mean(), (self,))
def backward():
self.grad += np.ones_like(self.data) * out.grad / self.data.size
out._backward = backward
return out
def backward(self):
order = []
visited = set()
def build(node):
if node in visited:
return
visited.add(node)
for parent in node.parents:
build(parent)
order.append(node)
build(self)
self.grad = np.ones_like(self.data)
for node in reversed(order):
node._backward()
# 用梯度下降拟合 y = 2x + 1。
x = Tensor([1.0, 2.0, 3.0, 4.0])
y = Tensor([3.0, 5.0, 7.0, 9.0])
w = Tensor(0.0)
b = Tensor(0.0)
for step in range(201):
w.grad.fill(0.0)
b.grad.fill(0.0)
prediction = w * x + b
error = prediction - y
loss = (error * error).mean()
loss.backward()
learning_rate = 0.05
w.data -= learning_rate * w.grad
b.data -= learning_rate * b.grad
if step % 50 == 0:
print(
f'step={step:3d} loss={loss.data:.6f} '
f'w={w.data:.4f} b={b.data:.4f}'
)
运行：
python tiny_autograd.py
训练结束时，w 应接近 2，b 应接近 1。这段代码展示了几个框架核心概念：
parents构成动态计算图。- 每个操作保存自己的局部导数规则。
backward()先拓扑排序，再逆序执行反向函数。- 广播不仅影响前向计算，也要求反向阶段把梯度恢复到参数原本的形状。
其中第四点很容易被简化实现忽略。前向里的标量偏置 b 被广播到整个向量，反向时则必须把每个样本对偏置的贡献累加起来。
从这个小引擎走到 Transformer
“从张量到 Transformer”不是一步跨越，而是一条依赖关系清晰的工程路径。可以按下面的顺序扩展：
1. 补齐张量算子
加入矩阵乘法、切片、转置、指数、对数和归约运算，并为每个算子实现反向规则。每增加一个操作，都应该做数值梯度检查：
# 用有限差分近似某个标量参数的梯度。
epsilon = 1e-6
numerical_grad = (loss_at(w + epsilon) - loss_at(w - epsilon)) / (2 * epsilon)
数值梯度不是训练方案，而是验证反向公式是否正确的测试工具。
2. 建立模块与参数系统
实现类似 Module、Parameter 和 Linear 的抽象，让模型能递归收集参数。这样优化器不需要了解网络结构，只需要遍历参数列表。
3. 实现稳定的神经网络组件
Softmax 不能直接对大数求指数，更稳妥的计算方式是先减去每行最大值：
shifted = logits - logits.max(axis=-1, keepdims=True)
probabilities = np.exp(shifted)
probabilities /= probabilities.sum(axis=-1, keepdims=True)
如果要把它纳入自己的自动微分框架，max、exp、除法和按维归约也必须具有相应的梯度规则。
4. 组合注意力与 Transformer 块
有了矩阵乘法、Softmax、线性层和归一化之后，才适合实现缩放点积注意力：
scores = (Q @ Kᵀ) / sqrt(head_dimension)
weights = softmax(scores + mask)
output = weights @ V
这时 Transformer 不再是神秘的大型组件，而是已有张量操作的组合。真正新增的难点会转向形状管理、掩码语义、数值稳定性以及内存成本。
教学框架不等于生产框架
亲手实现框架的目标是理解语义，而不是立刻替换 PyTorch。上面的示例就没有覆盖许多生产问题：
- 原地操作可能破坏反向传播所需的数据。
- 重复调用
backward()时，梯度累加语义需要明确定义。 - 广播、切片和复杂索引会让梯度传播更难处理。
- NumPy 实现无法自然获得 GPU 内核、算子融合和分布式执行能力。
- Transformer 训练还需要关注稳定初始化、掩码正确性和显存占用。
因此，更合理的使用方式是把 TinyTorch 一类项目当作“可执行的系统教材”：小规模实验在自制框架中完成，真实训练任务仍交给成熟生态。
建议的动手清单
如果准备沿这条路线学习，可以用以下检查表控制范围：
- [ ] 实现张量数据、形状和基础运算。
- [ ] 构建动态计算图与逆拓扑反向传播。
- [ ] 用有限差分检查每个新算子的梯度。
- [ ] 实现参数收集、线性层和 SGD。
- [ ] 训练一个线性回归或小型分类器。
- [ ] 加入稳定的 Softmax、归一化和交叉熵。
- [ ] 实现单头注意力，再扩展到多头注意力。
- [ ] 用小数据集验证 Transformer 块，而不是一开始追求大规模训练。
当你能解释一条梯度如何从损失穿过注意力、线性层和广播操作回到参数时，深度学习框架就不再只是一个需要导入的黑盒，而是一组可以拆解、测试和重新组合的系统机制。