计算图与算子：深度学习的基础构件

说实话，我刚入行那会儿，对计算图和算子这两个概念的理解特别模糊。总觉得框架都封装好了，我写个 conv2d 就完事了，管它底层怎么跑。直到有一次，我手写了一个自定义算子，结果反向传播死活不对，调试了整整两天...嗯，从那以后我才真正重视起这些基础概念。

计算图是什么？

计算图，说白了就是描述计算过程的有向无环图。每个节点代表一个操作（算子），每条边代表数据的流动（张量）。

我习惯把计算图想象成一张「菜谱」：

节点：就是每个烹饪步骤，比如「切菜」、「翻炒」
边：就是食材在步骤之间传递，比如「切好的土豆 → 下锅」

深度学习框架（PyTorch、TensorFlow）就是用这种方式来组织你的模型计算的。你写的每一行代码，框架都会默默构建一张计算图。

核心要点：计算图有两种模式——静态图和动态图。

静态图：先定义后执行（TensorFlow 1.x）。好处是优化空间大，但调试麻烦。
动态图：边定义边执行（PyTorch）。灵活，调试方便，我更喜欢这种。

我个人习惯用动态图，因为开发效率高。但如果你做推理部署，静态图往往能压榨出更多性能。

算子（Operator）的定义

算子就是计算图里的节点。每个算子完成一个特定的数学运算。

举个例子：

# 这是一个简单的计算图
# 算子：add（加法）
c = a + b

# 算子：mul（乘法）
d = c * e

# 算子：relu（激活函数）
f = relu(d)

这里 add、mul、relu 都是算子。每个算子都有：

输入：一个或多个张量
输出：一个或多个张量
属性：比如卷积的 stride、padding

我的经验：写算子时，一定要搞清楚输入输出的形状变化。我曾经因为一个 reshape 搞错了维度，导致整个模型训练 loss 直接炸了。排查了三个小时才发现是算子实现里少了一个转置。

前向与反向传播中的算子

这里有个关键点：每个算子其实对应两个函数——前向和反向。

前向传播：计算输出。比如卷积层，输入特征图，输出特征图。

反向传播：计算梯度。根据损失函数对输出的梯度，算出对输入的梯度，以及对参数（权重）的梯度。

你想想看，为什么框架能自动求梯度？就是因为每个算子都实现了自己的反向函数。框架沿着计算图反向走一遍，调用每个算子的反向函数，就能算出所有参数的梯度。

# 以矩阵乘法为例
# 前向：C = A @ B
# 反向：
#   dA = dC @ B.T
#   dB = A.T @ dC

注意：反向传播时，中间结果需要保存！比如上面的 A 和 B，反向计算梯度时需要用到。这就是为什么训练时显存占用比推理大得多——框架要缓存中间激活值。

我记得有一次优化一个 Transformer 模型，发现显存爆了。排查后发现是某个算子的反向实现里，缓存了一个巨大的中间张量。后来我手动重写了这个算子的反向逻辑，用重计算（recomputation）换显存，才把模型跑起来。

常见算子类型

咱们来看看深度学习里最常见的几种算子。这些你天天用，但你真的了解它们的计算模式吗？

1. Conv（卷积）

卷积是 CNN 的核心算子。它的计算模式是「滑动窗口 + 点积」。

我刚开始优化卷积时，直接写了个三重循环，结果慢得离谱。后来才知道，卷积可以转化为矩阵乘法（im2col + GEMM），或者用 Winograd 算法加速。

# 卷积的 naive 实现（千万别这么写！）
for n in range(N):
    for c in range(C_out):
        for h in range(H_out):
            for w in range(W_out):
                for kc in range(C_in):
                    for kh in range(K):
                        for kw in range(K):
                            output[n,c,h,w] += 
                                input[n,kc,h+kh,w+kw] * weight[c,kc,kh,kw]

优化思路：

im2col + GEMM：把卷积展开成矩阵乘法，利用 cuBLAS 加速
Winograd：减少乘法次数，适合小卷积核（3x3）
FFT：适合大卷积核

2. MatMul（矩阵乘法）

MatMul 是深度学习的基础操作。全连接层、注意力机制、甚至卷积优化都离不开它。

矩阵乘法的计算量是 O(n³)，是典型的计算密集型算子。优化 MatMul 的关键在于：

利用 GPU 的共享内存（tiling）
内存对齐，避免 bank conflict
使用 Tensor Core（如果硬件支持）

避坑指南：我曾经在写 MatMul 的 CUDA 实现时，忽略了矩阵的布局（行主序 vs 列主序）。结果算出来的结果全是错的。后来我养成了一个习惯——每次写矩阵运算前，先确认数据的内存布局。

3. Softmax

Softmax 常用于分类任务的输出层。它的计算很简单：

softmax(x_i) = exp(x_i) / sum(exp(x_j))

但这里有个数值稳定性问题。如果 x_i 很大，exp(x_i) 会溢出。所以实际实现中，我们会先减去最大值：

def softmax(x):
    # 减去最大值，防止溢出
    x_max = max(x)
    exp_x = [exp(v - x_max) for v in x]
    sum_exp = sum(exp_x)
    return [v / sum_exp for v in exp_x]

嗯，这里要注意：Softmax 的反向传播也很经典。它的梯度形式是：

dL/dx_i = s_i * (dL/ds_i - sum(dL/ds_j * s_j))

其中 s_i 是 softmax 的输出。这个公式在实现时，可以复用前向的结果，减少计算量。

4. 其他常见算子

算子类型	计算特点	优化要点
ReLU / GELU	逐元素操作，计算量小	内存带宽瓶颈，注意访存模式
BatchNorm / LayerNorm	需要计算均值和方差	减少全局同步，使用 Welford 算法
Embedding	查表操作	注意 embedding 矩阵的访存局部性
Attention	QKV 计算 + Softmax + 加权求和	Flash Attention 优化 IO

知识体系总览

下面这张图，是我自己总结的计算图与算子的知识体系。你可以把它当作一个「思维导图」来看：

这张图把本章的核心内容串起来了。从计算图的两种模式，到算子的定义和前后向传播，再到常见的算子类型。你每次写算子优化时，都可以回头看看这张图，确认自己当前在哪个环节。

总结一下：

计算图是框架组织计算的方式，分静态图和动态图
每个算子都有前向和反向两个函数
反向传播需要缓存中间结果，这是显存优化的关键
不同算子的计算模式和优化策略差异很大

好了，这一章的内容就到这里。记住，理解算子的底层实现，是做好算子优化的第一步。下一章我们会深入具体的优化技巧，到时候见。

📚 课程目录

计算图与算子：深度学习的基础构件

计算图是什么？

算子（Operator）的定义

前向与反向传播中的算子

常见算子类型

1. Conv（卷积）

2. MatMul（矩阵乘法）

3. Softmax

4. 其他常见算子

知识体系总览