• 文档
• 教程
• 人工智能

本篇导读：计算图（Computation Graph，也称数据流图 Data Flow Graph）是 AI 框架中最核心的数据结构，贯穿模型构建、自动微分、图优化、调度执行与部署的整个生命周期。本文将系统讲解计算图为何成为 AI 框架的统一抽象、其基本构成（张量与算子）、神经网络训练的五阶段流程，并梳理它与自动微分、控制流、图调度执行之间的深层关系。读完本文，你将建立起对主流 AI 框架（PyTorch、TensorFlow、MindSpore 等）前端核心技术的整体认知框架。

为什么 AI 框架需要计算图

在 AI 框架发展的最近一个阶段，技术上主要以计算图来描述神经网络。此前的实践最终催生出了工业级 AI 框架 TensorFlow 和 PyTorch，这一时期同时伴随着 Chainer、DyNet、CNTK、PaddlePaddle、JAX 等激发了框架设计灵感的诸多实验项目。

TensorFlow 和 PyTorch，特别是 PyTorch，代表了今天 AI 框架两种不同的设计路径：

• 系统性能优先改善灵活性：以 TensorFlow 为代表，优先考虑硬件亲和度与整体执行性能；
• 灵活性易用性优先改善系统性能：以 PyTorch 为代表，优先保证开发者的编程体验与调试便利。

这两种选择，随着神经网络算法研究和应用的更进一步发展，使得 AI 框架在技术实现方案上产生了巨大差异。

随着神经网络模型越来越复杂，包括混合专家模型（MOE）、生成对抗网络（GAN）、注意力模型（Attention Transformer）等，复杂的模型结构需要 AI 框架能够对模型算子的执行依赖关系、梯度计算以及训练参数进行快速高效的分析，以便：

• 优化模型结构；
• 制定调度执行策略；
• 实现自动化梯度计算；
• 提高 AI 框架训练的效率。

综上所述，目前主流的 AI 框架都选择使用计算图来抽象神经网络计算表达，通过通用的数据结构（张量）来理解、表达和执行神经网络模型，通过计算图可以把 AI 系统化的问题形象地表示出来。

AI 系统化问题的挑战

为了高效地训练一个复杂神经网络，AI 框架需要解决许多问题，例如：

• 如何对复杂的神经网络模型实现自动微分？
• 如何利用编译期的分析 Pass 对神经网络的具体执行计算进行化简、合并、变换？
• 如何规划基本计算 Kernel 在计算加速硬件（GPU/TPU/NPU）上高效执行？
• 如何将基本处理单元派发（Dispatch）到特定的高效后端实现？
• 如何对通过自动微分（反向传播实现）衍生的大量中间变量进行内存预分配和管理？

为了用统一的方式解决上述挑战，框架设计者需要为各类神经网络计算提供统一描述，从而在运行神经网络计算之前，能够对整个计算过程尽可能进行推断，在编译期自动为用户程序补全反向计算、规划执行，最大程度地降低运行时开销、复用和节省内存。这也正是计算图作为 AI 框架前端核心技术的重要原因。

计算图的定义：数据流图与计算图

我们经常遇到有些 AI 框架把统一的图描述称为数据流图，有些称为计算图，这里可以统称为计算图：

• 数据流图（Data Flow Diagram，DFD）：从数据传递和加工角度，以图形方式表达系统的逻辑功能、数据在系统内部的逻辑流向和逻辑变换过程。在 AI 框架中，数据流图表示对数据进行处理的单元：接收一定的数据输入，对其进行处理，再进行系统输出。
• 计算图（Computation Graph）：被定义为有向图，其中节点对应于数学运算，是表达和评估数学表达式的一种方式。在 AI 框架中，计算图就是一个表示运算的有向无环图（Directed Acyclic Graph，DAG）。

两者都将神经网络模型统一表示为图的形式，图由节点和边组成，都在描述数据在图中的节点间传播路径，数据在图中的传播过程就是对数据进行加工计算的过程。例如公式：

$$ f(x_1,x_2)=ln(x_1)+x_1x_2-sin(x_2) $$

就可以被转换为对应的计算图（详见计算图原理章节）。

计算图的基本构成：张量与算子

在 AI 框架中，计算图的基本组成有两个主要元素：

• 基本数据结构——张量（Tensor）：张量通过shape表示具体形状，决定内存中元素的组成结构与大小；其元素数据类型决定每个元素占用的字节数和实际内存空间大小。
• 基本计算单元——算子（Operator）：在加速器 GPU/NPU 中执行运算的是由最基本的代数算子组成，还会根据深度学习结构组成复杂算子。每个算子接受的输入输出不同，如 Conv 算子接受 3 个输入 Tensor、1 个输出 Tensor。

计算图描述的是运算的有向无环图，节点代表具体的计算操作（算子），边代表张量（数据流）。例如一个简单的卷积 + 激活的神经网络模型，其前向计算公式为：

$$ f(x) = ReLU(Conv(w, x, b)) $$

AI 框架中的计算图就会以 Conv 和 ReLU 为节点、以张量为边来表达这一计算逻辑，并在训练过程中由自动微分功能自动构建对应的反向计算图。

神经网络训练流程：五个阶段

神经网络的训练流程主要包括以下五个过程：

1. 前向计算：张量数据输入神经网络模型，模型按顺序（从输入层到输出层）计算并存储每层的结果，输出具体预测值；
2. 计算损失：根据损失函数（真实样本与模型预测的比较函数）对每个维度的参数求偏导数；
3. 自动求导：自动微分（Automatic Differentiation，AD）将复合函数分解为计算图，并以此计算任意两个节点间的梯度；
4. 反向传播：从损失函数（根节点）开始，沿计算路径向前追溯，逐节点计算误差项，将梯度传递给每一层参数；
5. 更新模型参数：根据优化器（Optimizer）的学习策略，如 $w = w - \eta \cdot grad$，小幅更新网络模型中的各个权重参数。

在基于计算图的 AI 框架中，这五个阶段统一表示为由基本算子构成的计算图，算子是数据流图中的节点，由后端进行高效实现。

其中，反向传播算法实际上就是自动微分：只要神经网络各个组件以及损失函数都是可微的，损失函数就是关于各输入变量的可微复合函数。自动微分把复合函数分解为输出变量（根节点）和一系列输入变量（叶子节点）及基本函数（中间节点），构成计算图，并依据加法法则（任意两节点间梯度为所有路径偏微分之和）与链式法则（一条路径的偏微分为各相邻节点间偏微分的连乘）计算梯度。AI 框架在实现自动微分时，对每个中间层存储的是向量-雅可比乘积（VJP）而非完整的雅可比矩阵，从而在减少存储的同时不影响导数计算。详细原理可参考计算图与自动微分。

计算图、控制流与调度执行

有了对计算图的基本了解之后，还需要深入理解计算图在真正执行计算时的两个关键问题。

控制流问题

在程序实现过程中，会遇到很多编程性问题，例如控制流（if、else、while、for 等）跟程序相关，而非跟计算和数学表示相关。计算图在数学上是 DAG，但编程语言中的分支与循环会引入环，从而无法进行有效的拓扑排序。目前基于计算图，控制流的解决方案主要有三类设计思路：

1. 复用宿主语言：以 PyTorch 为典型代表，用前端 Python 语言中的控制逻辑驱动后端计算图的执行；
2. 支持控制流原语：以 TensorFlow 为典型代表，后端原生支持 Switch、Merge、Enter、Exit、NextIteration 等控制流原语；
3. 源码解析：以 MindSpore 为典型代表，前端解析源码成计算图，后端将控制流解析成子图延伸计算图。

详细内容可参考计算图的控制流实现。

调度与执行方式

AI 框架根据计算图描述的数据依赖关系确定算子的执行顺序，由运行时系统调度计算图中的节点到设备上执行。计算图的执行方式分为两种模式：

1. 逐算子下发执行的交互式方式：如 PyTorch 框架；
2. 整个计算图或部分子图一次性下发到硬件执行：如 TensorFlow 和 MindSpore。

无论采用哪种模式，根据硬件能力的差异，异构计算图的执行又可分为三种模式：单算子执行（针对 CPU/GPU，逐个调度执行）、整图下沉执行（针对 DSA 架构 AI 芯片，如谷歌 TPU、华为昇腾 NPU，一次性下发整图减少 host-device 交互）、图切分到多设备执行（面向大模型，通过计算图切分与跨设备集合通信实现多设备并行）。同时，静态图在编译期可以充分利用计算图的全局信息进行算子融合、代数化简、内存复用等优化。详细内容可参考计算图的调度与执行。

动态计算图与静态计算图

根据反向计算图的构建时机，自动求导机制可分为基于对偶图、基于表达式追踪（Evaluation Trace）和基于图层 IR 三种方式：

• 动态计算图：每一次执行神经网络模型，依据前端语言描述动态生成一份临时的计算图，正向传播立即执行，反向传播后计算图立即销毁（PyTorch 的默认机制），易用性高、调试方便，但需要保存大量中间结果，内存占用较高；
• 静态计算图：根据 Python 等前端高级语言描述的神经网络拓扑结构和参数变量等图层信息，构建一个固定的计算图，执行前完成构图和编译优化，性能更高、便于全局优化与内存复用（TensorFlow 和 MindSpore 默认使用），但调试困难、控制流编写复杂。

动态图与静态图的转换又分为**基于追踪（Tracing）和基于源代码解析（Parsing）**两种方式，主流的 AI 框架正朝着动静统一的方向演进。详细内容可参考动态图与静态图转换。

使用计算图带来的好处

使用了计算图，能够方便 AI 框架以统一表示来描述神经网络训练的全过程。静态计算图可以在编译期对计算过程的数据依赖进行分析，带来以下好处：

• 简化数据流动过程：通过有向无环图的方式，避免无序的数据依赖；
• 动态和静态内存优化：神经网络模型执行中会产生固定和非固定的内存需求，基于计算图信息可提前对内存进行优化；
• 优化算子间的调度策略：方便进行算子的并行与调度执行优化；
• 改善运行时 Runtime 性能：通过计算图可以进行计算和通信的时间重叠优化。

通过静态计算图，AI 框架可以切分出三个解耦的优化层：计算图优化、运行时调度优化、算子/内核执行优化。针对新提出的神经网络模型结构和训练算法，扩展步骤也分为三阶段：在计算图层添加新算子定义；针对不同硬件内核实现计算优化；注册算子和内核函数，运行时派发硬件执行。

计算图的挑战与未来

面向未来，计算图还面临着若干演进方向，相关讨论详见计算图挑战与未来：

• 图表示：面对图神经网络（GNN）等大量稀疏结构，如何通过计算图更高效地计算稀疏图结构；
• 大数据融合：AI 框架的计算图如何与大数据处理（批处理、流处理、图数据处理等架构）进行融合与表示；
• 部署推理：既然最终形态是面向推理部署，AI 框架能否按照计算图设计后的逻辑明确分层解耦，泛化计算图解决部署流程复杂问题；
• 科学计算：面向多尺度科学计算问题，计算图能否足够表示科学计算的特殊数学范式。

学习路径建议

本篇为"计算图"系列的总览章节，建议按以下顺序深入阅读同一目录下的系列文档，它们共同构成 AI 框架前端核心技术的完整知识链：

1. 计算图原理：计算图定义、张量与算子基本构成、PyTorch 动态计算图与 Function 自定义；
2. 计算图与自动微分：前向计算、反向微分、雅可比矩阵与向量-雅可比乘积；
3. 计算图的调度与执行：算子调度（FIFO、并发、异构）、三种图执行模式与 PyTorch 算子两次调度；
4. 计算图的控制流实现：复用宿主语言、控制流原语与源码解析三种方案；
5. 动态图与静态图转换：动静分离、动静结合到动静统一的发展与 Tracing/Parsing 实现；
6. 计算图挑战与未来：图表示、大数据融合、部署推理与科学计算。

此外，本系列隶属于 AI 框架核心概述 课程，与自动微分基础、计算图调度等章节相互呼应。理解计算图这一贯穿 AI 框架整个生命周期的核心数据结构，是深入掌握 AI 编译器、推理引擎与分布式训练框架的必经之路。

• 文档
• 教程
• 人工智能