课程交接与内容概述#
这学期我来接替鹏帅老师,继续《人工智能中的编程》这门课程。我们的课程已经接近一半,好消息是,我们讲课的内容终于从 CUDA 回到了 Python 的范畴,难易程度相对要好一些。但同学们要完成课程 Lab,任务还是非常艰巨的。
这门课程旨在讲解 AI 相关编程框架所需了解的事情。之前鹏帅老师讲授了非常多硬核的 C++、CUDA 并行编程,这些是组成 AI 框架底层的各种重要算子,它们依赖硬件实现高速计算。接下来我们这部分,是 AI 框架中相对完整和核心的内容,主要研究 AI 框架本身,例如作为软件应完成哪些功能,重点将涉及计算图的构建以及自动微分等。
AI 系统可以大致分为三层:底层解决硬件快速计算问题;中间层解决面向用户快速搭建 AI 网络并使用的问题;系统层则是在 AI 能力扩展后,需要调度更多硬件进行统筹,涉及分布式系统。我们刚完成底层阶段,现在进入中间层。
回顾课程安排,我们已进行了八次 CUDA 相关的课程,大家已能实现各种算子,如矩阵乘、卷积和池化,这些都是 AI 框架中基础且常用的算法。后续我们关心的是如何面向用户。如果把 CUDA 直接暴露给用户,体验会比较痛苦。因此,我们通过混合编程提供接口,包装 CUDA 的复杂度,以便于 AI 训练的程序员使用更结构化的语言,便捷地构建 AI 计算图。
我们这部分总共也是八次课程,将讲解自动微分、计算图、编译器前端优化、编译器后端、数据处理以及异构处理器(CPU、GPU)的异步。最后还会有分布式相关的内容。
课程也贯穿着作业设计。大家已完成了前两个小作业。之前确认过,卷积池化这边确实是三次小作业的工作量。过渡到 Python 这一层,相应也会有两次小作业,一次是关于计算图和自动微分,这是我们这次课程就会开始讲的内容。考虑到有同学希望更好地安排时间,我们计划明天就把这个作业放出来。但大家不用着急,如果还在做 CUDA 的作业,我们预期的 deadline 还是会按正常节奏来,明天放的作业应该会有一个月或更充分的时间。相对来说,Python 层的这两个 Lab 工作量没有特别大,更多是希望大家了解代码背后的逻辑。
当然,同学们还一直带着一个大作业的任务,即未来手搓一个 PyTorch 1.0 水平的框架。Python 这一层写的计算图构建,有可能是需要直接迁移到自己框架中的。另外,计算图本身是一个代码填空的作业,提供的数据结构可能已包含 Python 上的 CPU 算子实现。大家需要想办法在自己的框架中,用之前 CUDA 实现的算子来完成计算图的构建。所以在做这两次作业时,大家要花时间构想一下,将来做大作业时,API 需要怎样设计,以及最开始用 CUDA 实现的算子是否有需要返工的地方。这两次作业会为大家积累经验和提供引导。
AI 框架的本质与演进#
要了解 AI 框架,我们首先要问一个很直觉的问题:AI 框架到底是什么?我们已经有具体的例子,比如 PyTorch,以及曾经流行但现在市场缩小的 TensorFlow。在后续课程中,我们会逐步讲解 AI 框架的设计与实现,以及它们如何演变成现在的市场分布。
课程最核心的问题是:为什么要有 AI 框架?它在 AI 编程中要完成哪些任务?我们要从深度学习的基础出发,推测 AI 框架应完成的工作。深度学习的训练涉及几个重要事情:神经网络的结构、数据驱动训练(数据和损失函数),以及基于梯度的优化。AI 框架需要服务于训练神经网络的全流程。
如果我们要有一个系统来服务深度学习,这个系统(AI 框架)大致需要完成几类任务。首先,希望用户能便捷地搭建网络,AI 框架作为一个隔离层,封装算子的复杂度,让用户通过 API 快速调用,如设计 Transformer 等复杂模型。其次,是微分的计算,即自动微分。接着,AI 框架需要兼具数据管理任务,处理大量的训练、测试、验证数据以及网络参数。最后,还包括各种优化、部署任务、面向硬件的加速任务以及分布式任务。
AI 框架所涉及的范围非常大,它的产生过程可以类比计算机操作系统。最初大家有各种单片机,后来发现需要一个软件来完成基础操作,让用户更便捷地使用,并打包主要功能。AI 领域任务越来越杂,使得我们迫切需要一个面向 AI 的管理系统,AI 框架因此诞生,它需要完成上述所有事情。
这些需求并非随着 AI 产生才被需要。AI 框架的设计借鉴了很多过去的经验。例如,数据处理软件积累了大量数据管理经验;硬件仿真或渲染领域积累了硬件使用经验(CUDA 最初就面向图形学);还有分布式系统的经验。除了这些现有工具可借鉴的需求外,AI 也带来了新的需求,主要是神经网络结构的搭建及相关优化工作。
在正式的 AI 框架出现之前,Caffe 和 Theano 这样的一些神经网络 Library 做了最早的尝试,但它们主要面向神经网络这部分新需求。当我们称之为一个 AI 框架时,我们期待它完成所有统筹工作:既包括神经网络的 Library 设计,也包括与数据相关的管理。AI 框架是过去工作(数据处理、分布式)和现有需求(神经网络、AutoDiff、硬件加速)的统一统筹管理。
在我们课程中,AI 框架的核心任务,如神经网络搭建、自动微分(AutoDiff)以及硬件操作,会是讲解的重点。数据管理等内容则介绍得相对较少。目前 AI 框架百花齐放,虽然 PyTorch 一家独大,但还有 Google 的 JAX、华为的 MindSpore、百度的 PaddlePaddle 等,它们背后有各自的设计、历史及市场因素。
AI 框架在整个学习工作中的定位,是隔离硬件和应用。上层是丰富的下游应用(如计算机视觉、大语言模型、机器人);底层是多元异构的硬件(如英伟达 GPU、国产昇腾、移动端 OpenCL)。AI 框架像操作系统一样,在中间做隔离,硬件只需对接框架,应用也只需对接框架,从而实现解耦,促进 AI 生态的发展。
AI 框架的层级结构与核心任务#
面向隔离硬件与应用的需求,AI 框架必然呈现层级结构。它有两大任务:一是面向应用,二是面向硬件。面向应用层,需要支持设计网络、完成训练,具体任务包括推导计算图、计算微分、网络优化、支持现有生态(如数据管理软件包),并提供简洁的编程语言(如 Python)和简便的模型搭建方式。面向硬件层,需要进行计算优化、适配硬件并行,并提供 API 让硬件驱动来对接。
AI 框架的所有任务构成了一个层次关系。最上层是编程接口,其任务是隔离编程难度,让用户使用 Python。框架根据 Python 代码进行后续处理,如构建计算图、执行自动微分,最终调用的可能是底层的 C 或 C++ 实现。
编程接口之下,是框架的核心表达,即计算图。计算图是神经网络训练中的一种表达和数据结构,我们下周会更正式地讲解。
有了计算图这个核心表达,就需要相应的处理操作,这催生了编译器前端。编译器前端的任务是根据编程语言(Python)自动构建计算图,并可能进行一些优化(如节点复用)。它之所以被称为编译器,是因为其需求、功能和设计与传统语言编译器(如 C++ 编译器)有诸多相似性。
构建并优化了计算图之后,需要逐步服务于硬件。这引出了编译器后端和运行时优化。编译器后端会面向特定硬件(比如 CPU 和 GPU 的不同配比)对计算图进行拆解和优化。这同样很像传统编译器调度计算机硬件以完成指令。
继续朝底层走,就是异构处理器(CPU、GPU)及其管理操作。
我们课程后续会重点讲解计算图的基础介绍。编程接口的转换过程(Python 到 C++)主要通过 Lab 让大家体会。课程还会重点讲解编译器前端,即如何构建和优化计算图。此外,面向 AI 服务的高效优化(自动微分)也是课程的重点。其他部分(如编译器后端、数据处理、分布式等)则会作为大概介绍。
自动微分(AutoDiff)的原理与应用#
我们已经了解 AI 框架更关心的是上层如何构建网络、构建计算图,以及如何通过计算图执行梯度计算。我们就从自动微分开始介绍 AI 框架。
在深度学习的计算中,核心是梯度下降,我们今后会持续关心这个抽象公式:
在这个公式里,最核心的计算就是微分梯度 的计算。因此,我们需要回顾一下求导的经典问题:过去有哪些实现手段?为什么它们不能满足 AI 的需求?自动微分为什么是天然适配的选择?
微分计算的传统方法与局限#
在有自动微分工具之前,微分计算有几种方式。
首先是手动微分 (Manual Differentiation)。以一个简单的递归表达式 为例,如果关心迭代四次的结果并对其求微分,手动推导的表达式会迅速膨胀,计算非常复杂,过程也非常痛苦。人手动计算微分显然不适用于复杂的 AI 编程。
其次是符号微分 (Symbolic Differentiation)。其核心逻辑是代替人做手动推导。它利用微分的链式法则和基本运算(加减乘除)的规则,自动化地构建微分表达式。其实现通常是:用户只写前向代码,符号微分软件(如 Python 的 sympy)记录所有原子操作,并以符号表达的形式管理这些数据结构。当需要求导时,软件根据推导规则构建出微分的函数表达式。这通常需要对操作符进行重载。
我们用 sympy 演示这个过程:使用 symbol 数据结构定义 ,然后编写前向循环计算 。调用 diff 函数得到符号表达的 df_dx。这个 df_dx 是一个字符串形式的符号表达式,我们可以打印并执行它。结果显示,即使是一个简单的前向操作,符号微分产生的反向微分表达式也可能“惨无人道”,极其膨胀,计算复杂度非常高。因此,符号微分也不能承担神经网络的微分计算任务。
最后是数值微分 (Numerical Differentiation)。这是在数值计算(如仿真)中常用的一种有误差的估计方法,最常用的是有限差分(Finite Difference)。例如,无偏的有限差分 。其优势在于,不需要关心反向微分的表达式,只要有前向计算的工具,就能估算微分。
但它有两个致命缺陷。第一个是效率问题。在神经网络中, 是网络的权重参数,数量极其庞大(可达几百亿)。数值微分需要对 每一个 标量参数 都调用一次(或两次)前向计算,这导致总计算量无法承受。
第二个是精度问题。 的选择非常有讲究。如果 选得太大,泰勒展开忽略的二阶误差项 本身就不是小量,会导致较大的截断误差 (Truncation Error)。如果 选得太小,又会遇到浮点数精度限制带来的舍入/抵消误差 (Rounding/Cancellation Error)。在神经网络中,(即损失函数 Loss)的输出范围很难控制,可能是一个绝对值非常大的数。当两个非常接近的大数 和 相减时,由于浮点数(特别是单精度 float)的精度不足,有效数字会相互抵消,导致巨大误差。因此, 常用 或 这样的值是在两种误差间权衡的结果。效率和精度的双重问题使得数值微分不适用于神经网络训练。
自动微分的核心机制:计算图#
鉴于传统微分方法的局限性,我们迎来了自动微分(AutoDiff)。自动微分更像符号微分,但有其特色。
自动微分的前提是,所有的计算都由有限的、我们知道如何求导的基本运算(如加减乘除、、、)构成。只要是这些基本运算的组合,我们就能通过链式法则(Chain Rule)进行计算。自动微分的核心思路就是使用计算图 (Computation Graph) 来迭代地运算链式法则。
以一个简单的多元函数 为例,我们可以构建其对应的计算图。这个计算图是一个有向无环图 (DAG),其中节点代表算符(操作符),边代表数据(变量)。计算图的构建通常在运行时通过操作符重载 (Operator Overloading) 实现。当代码执行如 +, *, ln 等操作时,框架会自动创建相应的节点和边,并按照运算符的优先级正确构建图的拓扑结构。
需要明确的是,这个计算图与神经网络的结构图不同。在神经网络的上下文中,计算图的输入 对应的是网络中需要求导的自变量,即网络的权重 (Weights);输出 对应的是因变量,即损失函数 (Loss)。由于损失函数通常是一个标量值,所以计算图是一个从高维(大量权重)输入到低维(单个 Loss)输出的巨大图结构。
自动微分为什么能避免符号微分的表达式爆炸问题?符号微分的目标是输出一个完整的、符号化的微分表达式。而自动微分(特别是反向模式)在图上遍历时,只计算并存储中间变量(节点)的数值梯度,例如 乘以 得到 。它只保留了这个最终的数值结果,而不是保留整个推导链条的符号表达式,从而避免了膨胀。
自动微分的两种模式:前向与反向#
自动微分依赖链式法则在计算图上进行遍历,这自然地产生了两种遍历方向:前向模式和反向模式。
前向模式 (Forward Mode),也称为切量模式 (Tangent Mode)。它从输入 出发,向输出 遍历。在遍历过程中,它维护一个称为切量 (Tangent) 的中间值 ,即中间节点对 某一个 输入 的导数。计算从 和其他 开始,按照图的拓扑顺序前向传播。例如,如果 ,则 。当遇到一个节点有多个输入时(例如 ),它的切量是输入的切量的线性组合()。
前向模式的特点是,一次前向遍历,可以计算出 所有 输出 相对于 某一个 输入 的导数。这对于神经网络(输入 维权重,输出 维 Loss)是极其低效的,因为它需要 次前向遍历才能计算出 Loss 对所有权重的梯度。
反向模式 (Reverse Mode),也称为伴随模式 (Adjoint Mode)。它从输出 出发,反向遍历到输入 。在遍历过程中,它维护一个称为伴随量 (Adjoint) 的中间值 ,即 某一个 输出 对中间节点的导数。计算从 开始,按照图的反向拓扑顺序传播。例如,如果 ,则 且 。
反向模式的关键在于处理旁路 (Bypass)。当一个节点 作为输入被多个后续节点(如 和 )使用时,它在反向传播中会从多条路径接收梯度。它的总伴随量 必须是所有这些路径梯度的总和,即 。
反向模式的特点是,一次反向遍历,可以计算出 某一个 输出 (即 Loss)相对于 所有 输入 (即 Weights)的导数。这完美匹配了神经网络梯度下降的需求( 个输入, 个输出),因此反向模式自动微分也被称为反向传播 (Backpropagation)。
从雅可比矩阵(Jacobian Matrix)( 维,代表输出 对输入 的导数)的角度看,前向模式(Tangent)计算的是 (雅可比-向量积),一次遍历得到 的一列信息;反向模式(Adjoint)计算的是 (向量-雅可比积),一次遍历得到 的一行信息。对于神经网络 的情况, 只有一行(即梯度向量),反向模式只需一次遍历即可获得。
反向模式的实现与扩展#
在实现反向模式时,为了处理旁路(一个节点有多个输出),我们需要存储从每条路径传回的部分伴随 (Partial Adjoint)。一个常见的实现方式是使用一个字典(例如 node_to_grad),键是节点,值是该节点收到的 Partial Adjoint 的列表或累加和。
反向模式的算法流程大致如下:首先对计算图进行反向拓扑排序。从输出节点 开始,其伴随量 。按照反向拓扑顺序遍历每个节点 :
- 从
node_to_grad中汇总 所有的 Partial Adjoint,得到其总伴随量 。 - 对于 的每一个输入 :
- 计算局部导数 。
- 根据链式法则计算传给 的 Partial Adjoint:。
- 将这个 存储到
node_to_grad[v_k]中。 遍历完成后,所有输入 节点的node_to_grad中存储的总和即为所需的梯度 。
在执行上述算法时,有两种选择:
- 历史计算 (Immediate Computation):在遍历时立即计算出所有伴随量的数值。这是早期框架(如 Caffe)采用的方式。
- 扩展计算图 (Graph Extension):在遍历时,并不立即计算数值,而是将反向传播的链式法则操作(乘法和加法)作为新的节点,构建出一个包含前向和反向计算的、更大的计算图。这是 TensorFlow 和 PyTorch 等现代框架采用的方式。
采用扩展计算图的方式具有显著优势。首先,构建的(前向+反向)图可以被复用 (Reusable)。对于固定的网络结构和损失函数,我们可以用同一张图反复计算不同输入数据下的梯度。
其次,它支持高阶微分 (Higher-Order Differentiation)。由于反向传播本身被表示为一张图,我们可以对这张“反向图”再次调用自动微分(即 grad 操作),从而实现微分的微分(如计算海森矩阵 Hessian)。采用历史计算的 Caffe 则无法支持此功能。但需要注意,高阶微分的图扩展可能是指数级的,会导致巨大的计算开销。
最重要的一点,扩展计算图使其可以被优化 (Optimizable)。一旦我们有了一个完整的计算图,就可以在执行计算之前,应用各种图优化算法(如节点合并、算子融合、化简),这正是 AI 编译器(前端和后端)的工作重点。这是现代框架选择图扩展模式的最主要原因。
自动微分的逻辑可以从标量无痛地扩展到向量和矩阵(张量)。我们只需要定义基础张量运算(如矩阵乘法 )的局部梯度(即它如何将伴随量反传给 和 )。只要定义了这些规则,它们就可以被插入到反向传播的链式法则中。我们 Lab 中实现的也将是基于矩阵的反向微分。
更有趣的是,只要我们能为某个操作定义清晰的前向行为和反向的伴随量传导规则,即便是非标准数学运算(例如从字典中取值 ),也可以被纳入自动微分框架(例如,定义 对 的贡献,而不影响 )。这极大地扩展了自动微分的能力,使其接近微分编程语言(如 JAX)。
关键点和注意事项#
- 在作业方面,卷积池化是三次小作业的工作量。Python 层的 Lab(计算图和自动微分)工作量相对不大,重在理解逻辑。
- 在做 Python Lab 时,要开始构想大作业(手搓 PyTorch 框架)的 API 设计,特别是思考如何将之前用 CUDA 实现的算子接入到新框架中。
- 符号微分(Symbolic Differentiation)的主要问题是会导致微分表达式爆炸性增长,计算复杂度高。
- 数值微分(Numerical Differentiation)有两个主要问题:一是效率低下,需要对每个参数单独计算前向;二是存在精度问题,受截断误差和舍入(抵消)误差的影响。
- 在 AI 框架的计算图中,输入 通常对应网络权重 (Weights),输出 对应损失函数 (Loss),Loss 通常是一个标量。
- 由于神经网络是多输入( 个权重)到单输出(1 个 Loss)的结构(),因此必须使用反向模式自动微分(Adjoint Mode / Backpropagation)。
- 前向模式(Tangent Mode)不适用于神经网络,因为它需要 次遍历才能获得所有梯度。
- 在实现反向模式时,当一个节点有多个输出(即在计算图中被多次使用)时,会产生旁路 (Bypass)。实现时必须正确汇总所有路径传回的部分伴随 (Partial Adjoint),例如使用字典结构来累加梯度。
- 现代 AI 框架(如 TensorFlow 和 PyTorch)采用扩展计算图 (Graph Extension) 模式,而不是立即计算(Immediate Computation)的模式(如 Caffe)。
- 采用扩展计算图的最主要原因是:图可以被优化。这引出了 AI 编译器的优化工作。
- 扩展计算图的其他优势包括:图可被复用,以及支持高阶微分(微分的微分)。但需注意高阶微分可能导致图的规模呈指数级增长。
- 自动微分可以从标量扩展到张量(矩阵)运算,只需为张量运算定义其局部的梯度(伴随量)传播规则。
- 自动微分甚至可以支持非数学运算(如字典查找),只要能清晰定义其前向行为和反向梯度传导规则。