【AI系统】计算图的控制流实现

本文涉及的产品
公共DNS(含HTTPDNS解析),每月1000万次HTTP解析
全局流量管理 GTM,标准版 1个月
云解析 DNS,旗舰版 1个月
简介: 计算图作为有向无环图(DAG),能够抽象神经网络模型,但在编程中遇到控制流语句(如if、else、while、for)时,如何表示成为难题。引入控制流后,开发者可构建更复杂的模型结构,但部署含控制流的模型至不支持Python的设备上较为困难。目前,PyTorch仅支持Python控制流,而TensorFlow通过引入控制流原语来解决此问题。计算图的动态与静态实现各有优劣,动态图易于调试,静态图利于优化。

计算图在数学上作为一个有向无环图(DAG,Directed Acyclic Graph),能够把神经网络模型的概念抽象出来作为同一描述,不过在计算机的编程中,会遇到很多 if、else、while、for 等控制流语句,有向无环图改如何表示控制流变成了计算图中一个很重要的问题。好处在于,引入控制流之后,开发者可以向计算图中引入分支选择以及循环控制逻辑,进而构造出更加复杂的神经网络模型结构。

目前,以 Pytorch 为例,它支持的仅仅是 Python Control Flow,即在 Python 层执行控制逻辑,而非计算图中支持控制流。这样就存在一个问题,如果要部署带 Control Flow 的模型就会比较困难,如何灵活部署带控制流的计算图到不支持 Python 的设备上?

计算图中的控制流实现,与控制流图并不是一个概念。在计算机科学中,控制流图 (CFG) 是程序执行期间所有可能路径的图形表示。控制流图是由 Frances E. Allen 发现的。他指出,Reese T. Prosser 之前曾使用布尔连接矩阵进行流分析。CFG 是许多编译器优化和静态分析工具不可或缺的一部分。

控制流

AI 框架作为一个可编程系统,在设计时一个首要设计选择是如何让开发者,能够独立于实现细节以最自然的方式描述出各类神经网络模型。描述的完备性不仅影响 AI 框架能所够支持的神经网络结构,决定了开发者在使用高级编程语言,去实现神经网络模型的灵活性,也影响 AI 框架后端优化的技术手段。

背景

在计算机科学中,控制流(Control Flow)定义了独立语句,指令,函数调用等执行或者求值的顺序。例如,根据函数 A 的输出值选择运行函数 B 或者 C 中的一个。

image

AI 框架把神经网络的计算过程,抽象为有向无环图。使用有向无环图描述神经网络计算的方式,符合算法开发者对神经网络的概念定义:算子间拓扑结构对学习特性有重要影响,可以通过计算图,方便地描述出大多数通过堆叠深度或多分枝形成的复杂神经网络。然而,随着神经网络算法的快速发展,一些新颖的神经网络结构很难自然地表示为纯计算图。

以 Transformer 结构的神经网络为例,来看看使用最自然地方式描述这些算法对 AI 框架会带来什么新的表示要求。Transformer 结构的神经网络算法中,图的左侧是一个通用 Transformer 结构的中关键步骤的计算示意图,通过堆叠 Transformer 结构使得网络模型层数更深,右侧对应了使用最自然的方式描述这一算法计算过程的伪代码。

image

Transformer 是一种基于注意力机制的神经网络结构,由多个 Encoder 和 Decoder 堆叠而成,可以应用于各种自然语言处理任务。在具体应用时,可以根据任务的特点和需求,选择不同的 Transformer 架构来构建模型。从伪代码描述中可以看到,想要以通用的方式,自然地描述出 Transformer 的算法框架,均依赖于循环控制逻辑 for

难点

引入控制流将会使得计算图的构建以及前向传播带来很大的差异。

首先,计算图将变为动态的方式,分支选择以及循环控制流只有在真实运行的时候,才能够依据其依赖的数据输入来判断走哪个分支、是否结束循环。

其次,控制流引入的另一个难点在于循环控制流的实现。引入循环之后,原本的计算图在逻辑上出现了环,从而无法进行有效的拓扑排序。所以对于有控制流的计算图,前向计算和反向传播的实现要么抛弃拓扑排序这一思路,要么就要通过其他手段将循环进行拆解。

为了能够支持含有控制流结构的神经网络模型,AI 框架需要引入了对动态控制流这一语言结构(language construct)的支持。目前基于计算图,在控制流解决方案上,主要采用了三类设计思路:

  • 复用宿主语言:复用前端宿主语言的控制流语言结构,用前端语言中的控制逻辑驱动后端计算图的执行;

  • 支持控制流原语:AI 框架的后端对控制流语言结构进行原生支持,计算图中允许计算流和控制流混合;

  • 源码解析:前端对高级编程语言的代码进行解析称计算图,后端对控制流语言结构解析成子图,对计算图进行延伸。

复用宿主语言以 PyTorch 为典型代表,支持控制流原语以 TensorFlow 为典型代表,源码解析的方式则以 MindSpore 为典型代表。

动态图

每一次执行神经网络模型,AI 框架会依据前端编程语言描述,动态生成一份临时的计算图(实际为单算子),这意味着该实现方式下计算图是动态生成,并且过程灵活可变,该特性有助于在神经网络结构调整阶段提高效率,这种实现方式也被称为动态计算图

复用宿主语言

PyTorch 采用的是动态图机制 (Dynamic Computational Graph),动态图使得 PyTorch 的调试变得十分简单,每一个步骤,每一个流程都可以被开发者精确的控制、调试、输出,甚至是在每个迭代都能够重构整个网络。

动态图中,通过复用宿主语言的控制流构建动态图,即复用 Python 等高级语言本身的控制流的执行方式。以 PyTorch 为代表,在实际执行计算的时候,当遇到 if、else、for 等控制流语句的时候,使用 Python 在 CPU 的原控制执行方式,当遇到使用 PyTorch 表示神经网络的 API 时,才继续按照动态计算图的执行方式进行计算。

下面的代码是 Transformer 结构对应的动态图示例 PyTorch 代码,Transformer 的 Decoder 结构中通过 for 循环,堆叠 Decoder 模块,最后通过 if 来检查是否需要加入 norm 正则化层。此时 AI 框架不再维护一个全局的神经网络描述,神经网络变成具体的 Python 代码,后端的张量计算以库的形式提供,维持了与 numpy 类似的一致编程接口。

from torch import nn

class TransformerDecoder(nn.Module):
    def __init__(self, decoder_layer, num_layers, norm=None):
        super().__init__()
        self.layers = _get_clones(decoder_layer, num_layers)
        self.num_layers = num_layers
        self.norm = norm

    def forward(self, tgt: Tensor, 
                memory: Tensor, 
                tgt_mask: Optional[Tensor] = None,
                memory_mask: Optional[Tensor] = None,
                tgt_key_padding_mask: Optional[Tensor] = None,
                memory_key_padding_mask: Optional[Tensor] = None,
                tgt_is_causal: Optional[bool] = None,
                memory_is_causal: bool = False):
        seq_len = _get_seq_len(tgt, self.layers[0].self_attn.batch_first)
        tgt_is_causal = _detect_is_causal_mask(tgt_mask, tgt_is_causal, seq_len)

        for mod in self.layers:
            output = mod(output, memory, tgt_mask=tgt_mask,
                         memory_mask=memory_mask,
                         tgt_key_padding_mask=tgt_key_padding_mask,
                         memory_key_padding_mask=memory_key_padding_mask,
                         tgt_is_causal=tgt_is_causal,
                         memory_is_causal=memory_is_causal)

        if self.norm is not None:
            output = self.norm(output)

        return output

代码出自 PyTorch 框架 github 官网 https://github.com/pytorch/pytorch/blob/main/torch/nn/modules/transformer.py

即使 PyTorch 2.X 版本出来,推出了图模式,能够生成真正的计算图,但是目前的方案中,遇到控制流,仍然会把网络模型切分成不同的子图来执行,遇到控制流会使用 Python 来执行调度。

复用宿主语言的方式,其优点在于:

  1. 由于用户能够自由地使用前端宿主语言 Python 代码中的控制流,即时输出张量计算的求值结果,有着更高的易用性;
  2. 模型即代码,动态图使用声明式编程的方式,使得定义神经网络模型的计算就像普通编写真正的程序。

这种复用宿主语言控制流语言驱动后端执行的方式有着更加友好的用户体验,但缺点也是明显的:

  1. 用户易于滥用前端语言特性,带来更复杂的性能问题;
  2. 计算图的执行流会在语言边界来回跳转,带来十分严重的运行时开销;
  3. 整体计算图把模型切分成按算子组成的小块,不利于神经网络模型大规模部署和应用;
  4. 控制流和数据流被严格地隔离在前端语言和后端编译优化之中,前端和后端执行不同的语言和编译体系,跨语言边界优化难。

静态图

AI 框架静态地生成可以根据 Python 等前端高级语言描述的神经网络拓扑结构,以及参数变量等图层信息构建一个固定的计算图,这种实现方式也称为静态计算图。具体在 AI 框架中,利用反向微分计算梯度通常实现为计算图上的一个优化 Pass,给定前向计算图,以损失函数为根节点广度优先遍历前向计算图的时,便能按照对偶结构自动生成出反向计算图。

控制流原语

支持控制流原语以 TensorFlow 为典型代表,在面对控制流需求时,计算图中引入控制流原语,运行时对控制流原语以第一等级(first-class)进行实现支持。其控制流的基本设计原则是:引入包含少量操作的原子操作符(即 Function 类特殊算子),在这些操作符之上来表达 TensorFlow 应用的复杂控制流。

TensorFlow 计算图中支持控制流的方案,主要分为 3 层。暴露给开发者用于构建计算图的前端 API,这些 API 会被转换成更低等级的控制流原语,再由计算图优化器进一步进行改写。为了平衡编程的易用性和优化器设计中保留更多易于被识别出的优化机会,TensorFlow 提供了多套有着不同抽象等级的前端 API 以及计算图上的控制流原语。

为了提高可理解性和编程效率避免开发者直接操作底层算子,这些计算图中的控制流原语会被封装为前端的控制流 API,下图是用户使用前端基础控制流 API 编写带条件和循环的计算,以及它们所对应的计算图表示。为了简化开发者识别计算图中的控制结构,TensorFlow 基于底层控制流原语,引入高层 Functional 控制流算子,同时添加高层控制流算子向底层控制流算子的转换逻辑。

image

下面以循环嵌套两个 for 循环代码,使用 TensorFlow 2.X 的 API 为例子:

i = tf.constant(0)
j = tf.constant(0)
a = lambda i: tf.less(i, 10)
b = lambda i: (tf.add(i, 1), )
c = lambda i: (tf.add(j, 1), )
y = tf.while_loop(a, b, [j])
r = tf.while_loop(c, y, [i])

TensorFlow 的计算图,每个算子的执行都位于一个执行帧中(execution frame)中,每个执行帧具有全局唯一的名字作为标识符,控制流原语负责创建和管理这些执行帧。可以将执行帧类比为程序语言中的域(Scope),其中通过 key-value 表保存着执行算子所需的上下文信息,如输入输出变量存储位置等。当计算图中引入控制流后,每个算子有可能被多次执行,控制流原语会在运行时创建这些执行帧,执行帧可以嵌套,对应了开发者写出的嵌套控制流。

# 具有全局唯一的名字作为标识符
execution frame:{
    # 保存着执行算子所需的上下文信息
     key(ops_1): value(input_addr, output_addr, ops_attr),
     key(ops_2): value(input_addr, output_addr, ops_attr),
     ...
     # 嵌套 Execution frame,可并发优化
     execution frame:{
     key(ops_1): value(input_addr, output_addr, ops_attr),
     key(ops_2): value(input_addr, output_addr, ops_attr),
     ...
    }
}

tf.while_loop的循环体是一个用户自定义计算子图,对于每个 while 循环,TensorFlow 运行时会设置一个执行帧,并在执行帧内运行 while 循环的所有操作。执行帧可以嵌套。嵌套的 while 循环在嵌套的执行帧中运行。位于同一个计算帧中,嵌套的tf.while_loop对应嵌套的计算帧,位于不同计算帧中的算子,只要它们之间不存在数据依赖,有能够被运行时调度并发执行。只要执行帧之间没有数据依赖关系,则来自不同执行帧的操作可以并行运行。

如下图所示,TensorFlow 的原子操作集之中有五个控制流原语运算符,其中 Switch 和 Merge 组合起来可以实现条件控制。所有五个基元一起组合则可以实现 while 循环。

image

其中:

  • Switch:Switch 运算符会根据输入控制张量 p 的布尔值,将输入张量 d 转发到两个输入中的一个。只有两个输入都准备好之后,Switch 操作才会执行。

  • Merge:Merge 运算符将其可用的输入之一转发到其输出。只要它的任何一个输入可用,merge 运算符就会执行。如果有多个可用的输入,则无法确定它的输出。

  • Enter(name):Enter 操作符将其输入转发到由给定名称唯一标识的执行帧。这个 Enter 操作用于将一个执行帧中的张量传递给一个子执行帧。对于同一个子执行帧可以有多个 Enter 操作,每个操作都会使子执行帧中的张量可用(异步)。当输入可用时,Enter 操作将执行。一个新的执行帧在执行该帧第一个 Enter 操作时候被实例化。

  • Exit:Exit 操作符将一个张量从一个执行帧返回给它的父执行帧。一个执行帧可以有多个 Exit 操作返回到父执行帧,每个操作都异步地将张量传回给父帧。当一个 Exit 的输入可用时,该 Exit 操作就被启用。

  • NextIteration: 一个 NextIteration 操作符将其输入转发到当前执行帧的下一个迭代。TensorFlow 运行时会跟踪维护执行帧中的迭代信息。一个执行帧中执行的任何操作都有一个唯一的迭代 ID,这使得我们能够唯一地识别迭代计算中同一操作的不同调用(比如 hile 操作之中,某一个 op 可能会多次执行)。

优点在于:

  1. 静态图由于能够在执行具体计算之前得到神经网络模型计算的全过程统一描述,使得 AI 编译器在编译优化期间能够利用计算图的信息进行优化

  2. 执行逻辑无需在前端宿主语言与运行时之间反复切换,因此往往有着更高的执行效率。

不过这种方式也有其比较大的缺点

  1. 由于控制流原语的语义设计,首要提升运行时的并发数和高效执行,与开发者在描述神经网络模型时候使用 Python 语法在编程习惯上差异较大,对开发者来说存在一定的易用性困扰

  2. 为了解决让开发者对细节控制原语不感知,需要对控制流原语进行再次封装,以控制流 API 的方式对外提供,这也导致了构建计算图步骤相对复杂。

源码解析

源码解析的方式是:前端对高级编程语言的代码进行解析称计算图,后端对控制流语言结构解析成子图,对计算图进行延伸。具体实现的过程中,计算图对能够表达的控制直接展开,如 for 循环内部的内容,直接展开成带顺序的多个计算子图。另外通过创建子图进行表示,运行时时候动态选择子图执行,如遇到 if 和 else 的分支时候分别创建 2 张子图存放在内存,当 DSA 的控制模块判断需要执行 if 分支的时候,把通过 if 分支产生的子图产生的序列调度进入执行队列中。

通过对高级语言的源码解析成计算图,在对计算图进行展开的方式,其优点在于:

  1. 用户能够一定程度自由地使用前端宿主的控制流语言,即在带有约束的前提下使用部分 Python 代码;

  2. 解耦宿主语言与执行过程,加速运行时执行效率;

  3. 计算图在编译期得到全计算过程描述,发掘运行时效率提升点;

因为属于静态图的方式,因此继承了声明式编程的缺点

  1. 硬件不支持的控制流方式下,执行流仍然会在不同编程语言的边界来回跳转,带来运行时开销;

  2. 部分宿主的控制流语言不能表示,带有一定约束性;

目录
相关文章
|
机器学习/深度学习 人工智能 自然语言处理
aiXcoder XL 智能编程大模型发布:自然语言一键生成方法级代码
aiXcoder XL 的出现,为程序员提供了大模型时代的个性化智能编程体验。随着 AI 技术的发展和普及,这或将重新定义编程领域。
742 0
aiXcoder XL 智能编程大模型发布:自然语言一键生成方法级代码
|
2月前
|
机器学习/深度学习 人工智能 算法
神经形态计算:模拟大脑的计算方式
【10月更文挑战第11天】神经形态计算作为一种新兴的计算范式,正以其独特的优势和广阔的应用前景吸引着越来越多的关注。通过模拟大脑的计算方式,神经形态计算不仅能够提高计算速度和能效,还能在处理复杂任务时展现出更高的智能水平。我们有理由相信,在未来的发展中,神经形态计算将为我们带来更多的惊喜和突破,引领我们进入一个全新的计算时代。
|
3天前
|
存储 人工智能 算法
【AI系统】计算图的优化策略
本文深入探讨了计算图的优化策略,包括算子替换、数据类型转换、存储优化等,旨在提升模型性能和资源利用效率。特别介绍了Flash Attention算法,通过分块计算和重算策略优化Transformer模型的注意力机制,显著减少了内存访问次数,提升了计算效率。此外,文章还讨论了内存优化技术,如Inplace operation和Memory sharing,进一步减少内存消耗,提高计算性能。
53 34
【AI系统】计算图的优化策略
|
3天前
|
机器学习/深度学习 存储 人工智能
【AI系统】卷积操作原理
本文详细介绍了卷积的数学原理及其在卷积神经网络(CNN)中的应用。卷积作为一种特殊的线性运算,是CNN处理图像任务的核心。文章从卷积的数学定义出发,通过信号处理的例子解释了卷积的过程,随后介绍了CNN中卷积计算的细节,包括卷积核、步长、填充等概念。文中还探讨了卷积的物理意义、性质及优化手段,如张量运算和内存布局优化。最后,提供了基于PyTorch的卷积实现示例,帮助读者理解和实现卷积计算。
63 31
【AI系统】卷积操作原理
|
2天前
|
机器学习/深度学习 存储 人工智能
【AI系统】计算图与自动微分
自动求导利用链式法则计算雅可比矩阵,从结果节点逆向追溯计算路径,适用于神经网络训练中损失值对网络参数的梯度计算。AI框架中,自动微分与反向传播紧密相连,通过构建计算图实现高效梯度计算,支持动态和静态计算图两种模式。动态图如PyTorch,适合灵活调试;静态图如TensorFlow,利于性能优化。
23 6
【AI系统】计算图与自动微分
|
2天前
|
机器学习/深度学习 人工智能 PyTorch
【AI系统】计算图基本介绍
近年来,AI框架如TensorFlow和PyTorch通过计算图描述神经网络,推动了AI技术的发展。计算图不仅抽象了神经网络的计算表达,还支持了模型算子的高效执行、梯度计算及参数训练。随着模型复杂度增加,如MOE、GAN、Attention Transformer等,AI框架需具备快速分析模型结构的能力,以优化训练效率。计算图与自动微分紧密结合,实现了从前向计算到反向传播的全流程自动化。
21 4
【AI系统】计算图基本介绍
|
2天前
|
机器学习/深度学习 人工智能 算法
【AI系统】计算图挑战与未来
当前主流AI框架采用计算图抽象神经网络计算,以张量和算子为核心元素,有效表达模型计算逻辑。计算图不仅简化数据流动,支持内存优化和算子调度,还促进了自动微分功能的实现,区分静态图和动态图两种形式。未来,计算图将在图神经网络、大数据融合、推理部署及科学计算等领域持续演进,适应更复杂的计算需求。
29 5
【AI系统】计算图挑战与未来
|
3天前
|
存储 机器学习/深度学习 人工智能
【AI系统】微分计算模式
本文深入探讨了自动微分技术,这是AI框架中的核心功能。自动微分分为前向微分和后向微分两种模式,主要通过雅克比矩阵实现。前向模式适用于输出维度大于输入的情况,而后向模式则更适合多参数场景,广泛应用于现代AI框架中。文章还详细解释了这两种模式的工作原理、优缺点及应用场景。
17 2
【AI系统】微分计算模式
|
2天前
|
机器学习/深度学习 人工智能 PyTorch
【AI系统】计算图原理
本文介绍了AI框架中使用计算图来抽象神经网络计算的必要性和优势,探讨了计算图的基本构成,包括标量、向量、矩阵、张量等数据结构及其操作,并详细解释了计算图如何帮助解决AI工程化中的挑战。此外,文章还通过PyTorch实例展示了动态计算图的特点和实现方法,包括节点(张量或函数)和边(依赖关系)的定义,以及如何通过自定义Function实现正向和反向传播逻辑。
26 7
|
2天前
|
人工智能 调度 算法框架/工具
【AI系统】计算图的调度与执行
深度学习训练过程涉及前向计算、计算损失及更新权重参数。AI框架通过计算图统一表示训练过程,算子作为计算图的节点,由后端硬件高效执行。计算图调度包括算子间的调度、并发调度和异构调度,确保计算资源的有效利用。图执行模式分为单算子执行、整图下沉执行和图切分多设备执行,适应不同场景需求。以PyTorch为例,其算子执行通过两次调度选择合适的Kernel进行张量操作,并支持自动求导。
22 5