第二章 深度学习基础：神经网络基本结构、张量运算、自动微分原理与常见模型的计算图表示

作为AI编译工程师，理解深度学习模型的计算本质是设计高效编译器、优化算子调度与内存分配的前提。本章将围绕神经网络的基本结构、张量运算的底层语义、自动微分的实现机制，以及CNN、RNN、Transformer三类主流模型的计算图表示展开，帮助读者建立从数学表达式到计算图IR（中间表示）的映射能力。

计算图 IR (Graph / SSA / MLIR) 神经网络结构 线性层 / 激活 / 损失 张量运算 MatMul / Conv / Broadcast 自动微分 反向模式 / 梯度累加 CNN Conv / Pool / FC RNN LSTM / GRU / 时序展开 Transformer Attention / FFN / LayerNorm 编译优化目标 算子融合 / 内存规划 / 并行调度 / 硬件映射

上图展示了本章的核心知识框架：神经网络结构定义了计算图的拓扑，张量运算构成了计算图的节点语义，自动微分则规定了梯度信息的反向传播路径。CNN、RNN、Transformer作为三种典型模型，其计算图表示分别呈现出不同的特征——CNN以空间卷积和池化为主，RNN引入时间循环依赖，Transformer则通过自注意力机制实现全局信息交互。AI编译工程师需要从这些高层抽象中提取出统一的IR表示，进而进行算子映射与优化。

2.1 神经网络基本结构：从神经元到计算图

核心原理

神经网络的基本组成单元是神经元，其数学表达为 \( y = f(Wx + b) \)，其中 \( W \) 为权重矩阵，\( b \) 为偏置，\( f \) 为非线性激活函数。在计算图视角下，每个神经元对应一个"线性变换 + 激活"的子图。现代深度学习框架（如PyTorch、TensorFlow）将网络表示为有向无环图（DAG），节点代表算子（Op），边代表张量（Tensor）的数据依赖关系。

从编译角度，神经网络结构可拆解为三个层次：模型层（定义网络拓扑）、算子层（具体计算逻辑）、内核层（硬件上的实际执行代码）。AI编译器的主要工作之一，就是将模型层的描述逐步lower到内核层，同时保持数学语义不变。

层次表示形式编译关注点
模型层Python类 / 网络结构描述拓扑解析、参数绑定
算子层计算图中的Op节点算子语义、shape推断、数据类型
内核层CUDA kernel / CPU SIMD代码循环优化、内存访问模式、向量化

判断方法或操作流程

在AI编译器中，解析神经网络结构通常遵循以下流程：

  1. 前端解析：读取模型文件（如ONNX、TorchScript），构建初始计算图。
  2. 规范化：将不同框架的算子统一映射到编译器内部定义的方言（dialect）上。
  3. Shape推断：沿数据流方向传播张量维度信息，标记动态维度。
  4. 控制流处理：识别循环、条件分支，转换为结构化控制流IR。
经验提醒：在解析RNN或动态shape模型时，计算图可能包含循环依赖。编译器需要将循环展开（unroll）或保留为循环IR，避免在编译期无限递归。

案例拆解：一个两层MLP的编译解析

假设有一个简单的两层全连接网络：输入维度784，隐藏层128，输出10。在编译器中，该网络被表示为：

%1 = matmul(%input, %W1)   // [batch, 784] x [784, 128] -> [batch, 128]
%2 = add(%1, %b1)          // 广播加法
%3 = relu(%2)              // 激活
%4 = matmul(%3, %W2)       // [batch, 128] x [128, 10] -> [batch, 10]
%5 = add(%4, %b2)
%6 = softmax(%5)           // 输出概率

编译器需要识别出matmul+add可以融合为带偏置的矩阵乘（如GEMM with bias），relu可以融合到前一个算子的epilogue中，从而减少内核启动次数和内存读写。

常见误区或总结复盘

风险提示：初学者常将"网络结构"与"计算图"混为一谈。网络结构是静态描述，而计算图是实例化后的数据流。同一网络结构在不同batch size下会生成不同shape的计算图，编译器必须处理这种动态性。

2.2 张量运算：语义、布局与广播规则

核心原理

张量（Tensor）是多维数组的泛化，其运算包括逐元素运算（add、mul）、归约运算（sum、max）、矩阵乘法（matmul）、卷积（conv2d）等。从编译视角，张量运算的关键属性包括：shape（维度信息）、strides（内存步长）、dtype（数据类型）、layout（内存布局，如NCHW vs NHWC）。

广播（Broadcasting）是张量运算中容易出错但极其重要的语义。当两个张量shape不完全一致时，NumPy和深度学习框架会按照从右向左对齐的规则自动扩展维度。编译器必须显式处理广播，否则会产生错误的内存访问。

判断方法或操作流程

在编译器中处理张量运算时，需要执行以下步骤：

  • Shape对齐：检查两个输入张量的维度是否满足广播条件（对应维度相等或其中一个为1）。
  • 内存布局转换：如果输入layout不一致（如一个NCHW、一个NHWC），插入layout转换算子。
  • 向量化分析：判断运算是否可向量化（SIMD），确定合适的向量宽度。
  • 内存分配：根据输出shape和layout计算所需内存大小及对齐要求。

案例拆解：广播语义在编译器中的实现

考虑计算 \( A[B, 1, D] + B[1, C, D] \)，输出shape为 \( [B, C, D] \)。在底层实现中，编译器不能简单地对两个张量做逐元素加法，因为内存访问模式不同。一种高效策略是：

// 伪代码：广播加法内核
for (b = 0; b < B; b++)
  for (c = 0; c < C; c++)
    for (d = 0; d < D; d++)
      out[b][c][d] = A[b][0][d] + B[0][c][d];

编译器通过循环重排和寄存器缓存，可以避免对A和B的重复加载。在AI加速器（如TPU）上，广播运算通常由硬件单元直接支持，编译器只需生成正确的地址映射。

常见误区或总结复盘

风险提示：广播虽然方便，但会引入隐式的数据复制。在内存受限的嵌入式设备上，编译器应尽量消除广播，通过修改循环顺序或使用strides来避免实际复制。

2.3 自动微分原理：反向模式与梯度图

核心原理

自动微分（Automatic Differentiation, AD）是深度学习训练的数学基础。与数值微分和符号微分不同，AD通过链式法则在计算图上精确计算导数。反向模式（Reverse Mode）是深度学习中最常用的模式：前向传播计算输出，反向传播从损失函数开始，逐层计算梯度并累加。

从编译角度，自动微分相当于对原始计算图进行变换，生成一个"反向计算图"。这个反向图包含梯度算子（如MatMul的梯度、ReLU的梯度）以及梯度累加逻辑。编译器需要决定何时进行梯度计算（即时vs延迟），以及如何复用前向传播的中间结果。

判断方法或操作流程

实现反向模式自动微分的关键步骤：

  1. 前向图记录：在算子执行时记录操作历史和中间张量。
  2. 拓扑排序：对前向图进行逆拓扑排序，确定反向传播顺序。
  3. 梯度算子生成：为每个前向算子生成对应的梯度算子（如add的梯度是split，matmul的梯度是额外的matmul）。
  4. 梯度累加：当某个张量被多个算子使用时，其梯度需要累加。

案例拆解：标量损失函数的梯度图

考虑一个简单的计算：\( loss = (x \cdot w - y)^2 \)。其前向计算图为：

%1 = mul(x, w)      // 逐元素乘
%2 = sub(%1, y)     // 减法
%3 = square(%2)     // 平方
%4 = reduce_mean(%3) // 标量损失

反向传播时，编译器生成：

%g4 = constant(1.0)                    // loss的梯度
%g3 = broadcast(%g4)                   // 广播到与%3同shape
%g2 = mul(%g3, 2 * %2)                 // square的梯度
%g1 = identity(%g2)                    // sub的梯度传递
%gw = mul(%g1, x)                      // 对w的梯度
%gx = mul(%g1, w)                      // 对x的梯度

编译器可以优化上述过程：将 \( 2 * %2 \) 融合到square的梯度中，减少一次乘法。

常见误区或总结复盘

经验提醒：梯度图通常比前向图大2-3倍，内存占用是训练编译优化的核心瓶颈。编译器可采用"重计算"（checkpointing）策略，在前向时丢弃部分中间结果，反向时重新计算，以内存换计算。

2.4 常见模型的计算图表示：CNN、RNN与Transformer

核心原理

三种主流模型的计算图各有特点，对编译器提出了不同挑战：

  • CNN：计算图以卷积、池化、全连接为主，结构规整，易于并行。但卷积运算在HW上的映射复杂（im2col、Winograd等算法选择）。
  • RNN：计算图包含时间维度的循环依赖，难以并行。编译器需要处理循环展开、序列长度动态变化等问题。
  • Transformer：以自注意力（Self-Attention）为核心，计算图包含大量的矩阵乘法、softmax和layer normalization。其计算密集且内存占用高，是编译器优化的重点对象。

判断方法或操作流程

针对不同模型，编译器需要采取不同的优化策略：

模型类型计算图特征编译优化重点
CNN静态shape、数据并行友好卷积算子选择、内存布局转换、算子融合
RNN动态时间步、循环依赖循环展开、序列打包、梯度裁剪
Transformer长序列、高内存占用Attention优化、KV Cache、张量并行

案例拆解：Transformer单层计算图的编译优化

以的单层解码器为例，其计算图包含：

// 输入 x: [batch, seq_len, hidden]
%1 = layer_norm(x)                    // 归一化
%2 = matmul(%1, Wq)                   // Query投影
%3 = matmul(%1, Wk)                   // Key投影
%4 = matmul(%1, Wv)                   // Value投影
%5 = reshape(%2, [b, h, s, d])        // 多头拆分
%6 = transpose(%5, [0, 2, 1, 3])      // 调整维度
%7 = matmul(%6, %6^T) / sqrt(d)       // 注意力分数
%8 = softmax(%7)                      // 归一化
%9 = matmul(%8, %6)                   // 加权求和
%10 = transpose(%9, [0, 2, 1, 3])
%11 = reshape(%10, [b, s, hidden])
%12 = matmul(%11, Wo)                 // 输出投影
%13 = add(%12, x)                     // 残差连接

编译器可以做的优化包括：将QKV三个matmul合并为一个大的batched matmul；将softmax的缩放因子融合到前面的matmul中；将layer_norm的均值/方差计算与归一化融合。在推理场景，还可以将KV缓存到专用内存，避免重复计算。

常见误区或总结复盘

风险提示：Transformer的动态序列长度（如不同batch的句子长度不同）会导致计算图shape不确定。编译器需要引入padding或使用动态shape支持，否则无法提前分配内存和调度内核。

总结本章，AI编译工程师必须深刻理解神经网络的结构语义、张量运算的底层行为、自动微分的图变换逻辑，以及典型模型的计算图特征。这些知识是后续学习算子优化、内存规划、并行调度等高级主题的基础。在实际工程中，建议读者多使用 torch.jit.script 或 onnxruntime 等工具导出计算图，观察其IR表示，逐步培养"图思维"。

← 上一章 📖 返回目录 下一章 →