第3章 训练与推理性能优化 随着深度学习的快速发展,为了提升神经网络的精度和泛化能力,数据集和参数量都在呈指数级向上攀升。超大规模网络的训练和推理,给AI框架带来了性能墙、效率墙及精度墙等多重挑战。 MindSpore框架通过技术创新及与昇腾处理器协同,提供了丰富的性能优化手段,实现了高效的运行态和部署态,大大提高了计算性能。 自动并行作为关键特性,实现了数据并行加模型的自动混合并行训练方式,降低了分布式训练难度,提高了算法研发效率,同时又能保持训练的高性能。二阶优化方法可以有效加速模型收敛,减少迭代次数。图算融合是MindSpore特有的网络性能优化技术,它通过分析和优化现有网络计算图逻辑,对原有计算逻辑进行拆分、重组、融合等操作,以减少算子执行间隙的开销并且提升计算资源利用率,从而实现网络整体执行时间的优化。 在部署态,除了要保证模型的精度不能下降,MindSpore还致力于减小推理模型的尺寸、降低模型推理时的时延。尤其是随着物联网概念的逐渐落地和端边设备的普及,在端边设备上部署AI框架对于高性能、低功耗、轻量化等推理指标有了更高的要求。 MindSpore提供模型量化、推理图优化、kenerl优化等关键技术,不仅减小了模型尺寸、内存占用,而且充分利用硬件计算资源,提升了推理性能,并通过常量折叠、低精度计算等技术降低了模型推理时的功耗。 本章将对上述技术逐一展开描述。介绍了训练性能优化技术: 3.1节介绍千亿参数模型的自动并行,3.2节介绍可以加速模型收敛的自研二阶优化器THOR,3.3节介绍感知量化训练和训练后量化两种模型压缩方法,3.4节和3.5节分别介绍网络性能优化技术中的类型推导和图算融合; 3.6节介绍算子融合和重排等推理图优化技术,3.7节则从算法和硬件加速两方面来探讨kenerl优化技术。 3.1千亿参数模型自动并行 近年来,越来越复杂的深度神经网络(Deep Neural Networks,DNNs)模型被构造出来,在各个应用中达到了创纪录的精度。这些成功的应用依赖于训练数据的规模、模型的参数量和硬件加速器的提升。训练数据从几百兆字节增长到数十太字节。自然语言处理(Natural Language Processing,NLP)领域的模型参数量从百万级别增长到千亿级别,计算机视觉领域的模型参数也经历了类似的增长。而主流的用于训练的加速器(GPU、TPU和Ascend)的内存仍然只有几十吉字节,因此为了能够使用大数据集训练大模型并且高效执行,机器学习框架需要提供丰富的并行训练能力。这里的“并行训练”包括数据并行、模型并行及混合并行等利用加速器集群进行分布式训练的方式。 本节首先介绍分布式训练相关的基础概念和关键问题,然后分别介绍MindSpore中用于并行训练的各个技术,最后通过GPT3千亿模型训练的实践介绍并行训练的多个技术的综合应用。 3.1.1分布式训练基础 1. 数据并行、模型并行和混合并行 简化版的线性神经网络和典型的并行训练时间线如图3.1所示。深度神经网络训练的目的是在给定训练数据集的前提下,找到合适参数(如图3.1(a)中的椭圆)来最小化损失函数(loss function)。通常的做法是使用迭代式算法和反向传播(backpropagation)来训练神经网络,直至损失函数值收敛到给定值。每次迭代包括两个阶段: 正向阶段会计算一批训练数据的当前预测值,该预测值将会用来计算损失值; 反向阶段会根据损失值计算参数的梯度,并利用梯度更新相应的参数值,下一次迭代会使用更新后的参数值。神经网络是由细粒度的算子(operator)构成的,且反向传播依赖于链式求导法则,所以反向阶段是由正向阶段的算子导数依次串成的。 图3.1简化版的线性神经网络和典型的并行训练时间线 当训练数据集过大时,应用最普遍的并行训练方式是数据并行: 每台设备具有全部参数,每台设备收到不同的训练数据(即切分数据而不切分参数),每次迭代即将完成时,每台设备会进行通信以同步计算出来的参数梯度值,然后利用同步后的梯度值更新参数,保证参数值在每台设备上的一致性。图3.1(b)为数据并行训练的时间线。模型并行是把参数切分到各个设备上,且在训练时每台设备的训练数据是相同的(即不切分数据而切分参数)。由于神经网络模型的参数会有多个,其可能的切分策略会不同,如图3.1(c)中的例子不会切分参数1(第一个矩阵乘法算子切分输入数据)而切分参数2,这样的并行方式叫作混合并行(即同时切分数据和参数)。 2. 集合通信、计算和通信交叠 集合通信是一组进程之间的通信方式,是由多个基本操作算子组成的,如图3.2所示的四种算子: 全汇聚(AllGather)、归并打散(ReduceScatter)、广播(Broadcast)和全归并(AllReduce)。AllReduce(AR)用于数据并行进程间同步梯度信息时,如图3.1(b)所示。AllGather(AG)和ReduceScatter(RS)分别用于进程间数据的汇聚和取切片的操作,如图3.1(c)所示。 注意到图3.1(b)、(c)中的AR操作是与反向梯度计算交叠进行的,这是由于后续梯度计算不依赖于此AR操作,交叠进行能够提升端到端的训练效率。 图3.2集合通信示意图 3.1.2关键问题 实现模型并行的方法有拆分算子、流水线(pipeline)并行、拆分并行子图等,优化内存使用的方法有切分优化器、异构并行、重计算等。想要实现模型训练端到端的性能最优,需要依赖上述多种技术的组合。因为模型结构各异,针对每个模型所选的技术组合也会不同。 如图3.3所示,目前基于Transformer结构的模型是最成功的语言模型,其中的代表模型是GPT3 (1750亿参数)。训练这类大模型,往往需要更大的数据集,所以就必须采用数据并行叠加模型并行的方式。在推荐领域,如Wide&Deep,特征数量是百亿甚至千亿规模,会有超大的头部嵌入(embedding)层,内存开销远超单加速器的内存,因此就必须使用模型并行将嵌入层参数切分到集群。此类模型采用模型并行转数据并行的方式是性能最高的。在人脸识别和超大规模图像分类领域,如ResNet类模型,分类数可能高达千万,尾部全连接层参数量能达到百亿级别。这类网络即使在数万分类的场景下,模型并行相对数据并行来说性能更优,因为数据并行通信的是参数梯度,通信数据量等于参数量,而模型并行通信的是特征映射(feature map),这类模型的特征映射数据量远小于参数的数据量,进而在通信上能够性能更优。 图3.3三种典型的模型结构 注: 图中矩形的高度代表参数量的大小。 如上所述,混合并行不同于数据并行,其有很多技术可选。现在主流框架的做法都是手动切分。然而,对于模型开发者来说,想要获得训练的高性能,有以下难点。 (1) 需要理解加速器集群的拓扑,以及节点内和节点间的带宽。一般地,服务器内多加速器间、机柜内服务器间、机柜间的带宽和延时依次降低。因此,需要把子模型间通信量大的放到节点内,通信量小的放到节点间。 (2) 需要考虑加速器的内存上限,在此约束下,找到性能较好的模型切分策略,该策略要保证切分后的各个子模型的计算量相对均衡。由于候选的并行方式巨大,从中选出一个最优的是难题。 (3) 手动切分的方式会显示地引入设备绑定和通信的代码,导致并行计算的逻辑与模型逻辑耦合在一起,增加了模型开发者的工作量。 3.1.3整体流程 神经网络模型可抽象为计算图,其中点表示算子和张量,边表示数据依赖关系。神经网络模型的并行训练可以抽象为把计算图切分映射到设备的切分,如图3.4所示。MindSpore并行训练的目标是构建一种易用高效的分布式并行训练模式,融合多种并行策略,让模型开发者不需关注使用哪种模式进行训练,具体策略如下: (1) 简化分布式并行编程,串行代码实现分布式训练,对用户屏蔽并行细节; (2) 实现上统一多种并行策略,一套框架支持多种并行策略任意组合; (3) 结合集群拓扑,保持训练高性能。 图3.4计算图切分到设备的切分映射关系 如图3.5所示,MindSpore中涉及四种类型的并行切分: 流水线切分、并行子图切分、算子切分和异构图切分。粗略地,流水线切分将计算图的线性结构切分到不同设备上,例如将图3.4中的算子5之前的算子切分到一设备,将剩余算子切分到另一设备; 并行子图切分将计算图中的可并行的部分切分到不同设备上,例如图3.4中的算子2、算子3、算子4是并行的,故可将其分配到不同设备上; 算子切分指的是将每个算子涉及的张量进行切分,进而使得每个算子的计算逻辑分配到每台设备上; 异构图切分指的是将一部分内存消耗巨大的算子放到主存(Host),利用CPU进行计算,另一部分放到加速器上执行。后续章节将会详细介绍这些切分方式。 注意到涉及图切分的每个模块都有待决策的变量,如算子切分模块中的每个张量的切分方法,综合起来的策略空间是巨大的。MindSpore中采用代价模型来评估每种给定并行策略的代价,并利用策略搜索算法来确定最优的并行策略。 图3.5MindSpore分布式并行相关流程(实线矩形代表对计算图进行切分的模块) 3.1.4流水线并行 流水线并行是将神经网络中的算子切分成多个阶段,再把不同阶段映射到不同的设备上,使得不同设备去计算神经网络的不同部分。流水线并行适用于模型是线性的图结构。如图3.6所示,将4层MatMul的网络切分成4个阶段,分布到4台设备上。正向计算时,每台机器在算完本台机器上的MatMul之后将结果通过通信算子发送给下一台机器,同时下一台机器通过通信算子接收上一台机器的MatMul结果,并开始计算本台机器上的MatMul; 反向计算时,最后一台机器的梯度算完之后,将结果发送给上一台机器,同时上一台机器接收最后一台机器的梯度结果,并开始计算本台机器的反向。 图3.6流水线并行的图切分示意图 简单地将模型切分到多个设备上并不会带来性能的提升,因为模型的线性结构导致同一时刻只有一台设备在工作,而其他设备在等待,造成了资源的浪费。为了提升效率,流水线并行进一步将小批次(minibatch)切分成更细粒度的微批次(microbatch),在微批次中采用流水线式的执行顺序,从而达到提升效率的目的,如图3.7所示。将小批次切分成4个微批次,4个微批次在4个组上执行形成流水线。微批次的梯度汇聚后用来更新参数,其中每台设备只存有并更新对应组的参数。序号代表微批次的索引。 图3.7带MicroBatch的流水线并行执行时间线示意图 MindSpore支持流水线并行,并对执行顺序进行了调整,来达到更优的内存管理。如图3.8所示,在第0 MicroBatch的正向执行完后立即执行其反向,这样做使得第0 MicroBatch的中间结果的内存得以更早地(相较于图3.7)释放,进而确保内存使用的峰值比图3.7的方式更低。 图3.8MindSpore流水线并行执行时间线示意图 MindSpore的流水线并行是以层为粒度进行配置的,代码3.1可实现简单的流水线并行: 即将matmul0算子分到第0阶段,将matmul1算子分到第1阶段。 代码3.1MindSpore中使用流水线并行示例 context.set_auto_parallel_context(parallel_mode="semi_auto_parallel", pipeline_stages=2) class MatMulCell(Cell): def __init__(self): super(MatMulCell, self).__init__() self.matmul = P.MatMul() def construct(self, x, y): out = self.matmul(x, y) return out class PipelineNet(Cell): def __init__(self): super(PipelineNet, self).__init__() self.matmul0 = MatMul() self.matmul0.pipeline_stage = 0 self.matmul1 = MatMul() self.matmul1.pipeline_stage = 1 def construct(self, x, y, z): out = self.matmul0(x, y) out = self.matmul1(out, z) return out 3.1.5并行子图切分 流水线并行适用于具有线性结构的模型,如果模型不具备这种结构呢?本节将介绍一种并行子图的切分方式。Transformer模型在NLP中有广泛的应用,但超大的模型参数量给训练所需要的资源带来了挑战。一种增大模型参数量,但保持训练的计算需求基本不变的方式是混合专家MoE (Mixture of Expert),即将Transformer中编码器(Encoder)的前馈网络(Feed Forward Network,FFN)替换成MoE结构。如图3.9(b)中的虚线框所示,经由路由门(Gate)计算后,训练数据会独立地分发到一个或多个FFN中,而后多个FFN在计算完成后汇聚成最后结果。这里FFN没有依赖关系,属于并行子图。Transformer中编码器扩展为MoE结构如图3.9所示。 图3.9Transformer中编码器扩展为MoE结构 如何高效地执行这种图结构?一种实现方式如下,将Gate函数的计算并行化,同时将多个FFN分配到每台设备上,每台设备得到一个或多个FFN,每台设备负责的参数量有效地下降了。若训练数据经过Gate计算后路由到本台设备负责的FFN,那么其直接传递给本台设备的FFN; 如果路由到其他设备的FFN,那么会经过全到全(AllToAll)通信,将训练数据发送到目的设备; 同样地,在经过FFN计算后,数据需再次经过AllToAll通信,把相应的结果路由回原设备。图3.10中展示了每台设备只有一个FFN的并行执行情况。 当每条训练数据仅会路由到较少FFN时,这种并行方式会非常有效,因为大大降低了内存使用,同时产生的通信量是较小的。 图3.10每台设备只有一个FFN的并行执行情况 3.1.6算子级并行 算子级并行即对神经网络中的每个算子进行切分,使得切分后的算子在每张卡上都进行计算,其汇总的结果是整个算子的计算结果。MindSpore支持算子级并行,对正向网络中的每个算子都独立建模。先看下面的例子,在该例中要计算两个连续的二维矩阵乘法(W和V是需要训练的参数,X是输入数据): Z=X×W×V (3.1) 假定当前总共有4个卡,第一个矩阵Y=X×W,用户想把X按行切4份(即数据并行); 而第二个矩阵Z=Y×V,用户想把V按列切4份(即模型并行): Y=X1 X2 X3 X4×W=Y1 Y2 Y3 Y4 (3.2) Z=Y×V1V2V3V4=Z1Z2Z3Z4 (3.3) 代码3.2即可实现上述的算子级并行计算: 代码3.2MindSpore中使用算子级并行示例 context.set_auto_parallel_context(parallel_mode="semi_auto_parallel", device_num=4) class DenseMatMulNet(nn.Cell): def __init__(self): super(DenseMutMulNet, self).__init__() self.matmul1 = ops.MatMul.shard(((4, 1), (1, 1))) self.matmul2 = ops.MatMul.shard(((1, 1), (1, 4))) def construct(self, x, w, v): y = self.matmul1(x, w) z = self.matmul2(y, v) return z 其中通过set_auto_parallel_context()接口将并行模式配置为半自动并行(用户可以手动配置每个算子的切分策略),并指定环境中的卡数; 通过shard接口来描述MatMul算子两个输入的切分策略,第一个括号内的数据配置表示算子的第一个输入张量的切分,第二个括号内的配置表示算子的第二个输入张量的切分,数值表示张量的对应维度被切分为几份。比如,ops.MatMul.shard(((4,1),(1,1)))表示对第一个输入张量的行切4份,列切1份(不切),对第二个输入张量的行切1份(不切),列切1份(不切)。 在算子级并行中有两个关键的建模: 分布式张量分布(tensor layout)和张量重排布(tensor redistribution)。张量分布表示张量切分后,张量切片在集群分布情况。张量重排布表示两种张量分布之间的转换。模型切分流程如图3.11所示,首先对每个算子的输入张量按策略进行切分,生成算子输入的张量分布; 然后根据算子的数学定义,推导出输出张量分布; 最后再检查前一个算子输出张量分布和下一个算子的输入张量分布,如果两种不同,则会插入张量重排布。 图3.11模型切分流程 张量分布表示一个完整张量切分到集群。张量可以各维度切分到集群,也可以在集群上复制。一个二维矩阵,切分到两台设备,有三种张量分布,行切分、列切分及复制,如图3.12所示。 图3.12二维矩阵切分到两台设备产生的三种张量分布示例 图3.13是一个更复杂的例子,是二维矩阵切分到四台设备。有四种张量分布,行列同时切分、复制、行切分+复制、列切分+复制。由于执行时间是计算时间加上通信时间,完全切分虽然计算时间更短,但是通信时间会比较长,有些场景下,不切满性能会更好,虽然计算时间更长,但是通信时间少,总的执行时间就会相对较短。 图3.13二维矩阵切分到四台设备产生的四种张量分布示例 张量重排布是不同张量分布之间的转换,张量在集群上从一种分布转成另外一种分布。在算子切分中,每个算子都是独立建模的,前一个算子输出的张量分布和下一个算子输入的张量分布可能会不同,编译阶段会自动插入重排布来调整输入输出关系。重排操作借助于MindSpore分布式通信原语,由“集合通信+split+concat”的算子组合。图3.14和图3.15说明了几种张量重排布的操作。 图3.14张量在两个节点的重排布 图3.15张量在四个节点的重排布 再回到之前两个矩阵相乘的例子,由于第一个算子输出的张量是第0维切分到集群,而第二个算子要求第一个输入张量在集群上复制,所以在图编译阶段,会自动识别两个算子输出与输入之间张量分布的不同,从而自动推导出所需的张量重排布的算子,自动在图中插入。而这个例子所需要的张量重排布就是一个AllGather算子(MindSpore AllGather算子会自动把多个输入张量在第0维合并成一个大的张量),如图3.16所示。 图3.16由数据并行转换到模型并行示例 3.1.7优化器切分 传统的数据并行模式将模型参数在每台设备上都有保有副本,把训练数据切分,在每次迭代后利用通信算子同步梯度信息,最后通过优化器计算对参数进行更新。数据并行虽然能够有效提升训练吞吐量,但并没有最大限度地利用机器资源。其中优化器会引入冗余内存和计算,消除这些冗余是本节关注的优化点。 在一个训练迭代中,数据并行为了收集各卡上不同样本产生的参数梯度,引入通信操作将梯度在多卡间同步。因为不涉及模型并行,每张卡上的优化器运算其实是基于相同的参数、在相同的方向上更新,而消除优化器冗余的根本思想就是将这部分内存和计算量分散到各个卡上,实现内存和性能的收益。 对优化器实现并行运算,有两种实现思路: 参数分组(weights grouping)和参数切分(weights sharding)。其中参数分组是将优化器内的参数及梯度做层间划分,大致的训练流程如图3.17所示,其中的FP指的是正向计算过程,BP指的是反向计算过程。将参数和梯度分组放到不同卡上更新,再通过通信广播操作在设备间共享更新后的权值。该方案的内存和性能收益取决于参数比例最大的参数分组。当参数均匀划分时,理论上的正收益是N-1/N的优化器运行时间和动态内存,以及N-1/N的优化器状态参数内存大小,其中N表示设备数,而引入的负收益是共享网络权重时带来的通信时间。 图3.17参数分组训练流程示意图 另一种实现参数切分的方式是对参数做层内划分,对每一个参数及梯度根据设备号取其对应切片,各自更新后再调用通信聚合操作在设备间共享参数。这种方案的优点是天然支持负载均衡,即每张卡上参数量和计算量一致,缺点是对参数形状有整除设备数要求。如果对网络中的权重也做切分,可以进一步减少静态内存,但这就需要将迭代末尾的共享权重操作移动到下一轮迭代的正向启动前执行,保证进入正反向运算的依旧是原始张量形状。此外,优化器并行运算带来的主要负收益是共享权重的通信时间,如果能够将其减少或隐藏,就可以带来性能上的提升。通信跨迭代执行的一个好处就是,可以通过对通信算子适当分组融合,将通信操作与正向网络交叠执行,从而尽可能隐藏通信耗时。通信耗时还与通信量有关,对于涉及混合精度的网络,如果能够使用fp16通信,通信量相比fp32将减少一半。综合上述特点,参数切分的实现方案如图3.18所示。 图3.18参数切分的实现方案 在实际网络训练的测试验证中,参数切分带来的内存收益是显著的。尤其是对于大规模网络模型而言,通常选择当下流行的Adaptive Moment estimation (Adam)和Layerwise Adaptive Moments optimizer for Batching training (LAMB)训练网络,优化器自身的参数量和计算量不容忽视。经过参数分组,网络中的权重参数和优化器中的两份状态参数都减少了N-1/N倍,极大地节省了静态内存空间。这为增大单轮迭代样本数量、提升整体训练吞吐量提供了可能,有效解决了大规模网络训练的内存压力。 除前面提到的跨迭代通信并行和混合精度加速外,MindSpore实现的优化器参数切分还具有与算子级模型并行混合使用的优势。当算子级模型并行参数未切满时,可以继续在数据并行的维度上进行优化器参数切分,增大机器资源的利用率,从而提升端到端性能。优化器参数切分配置如代码3.3所示。 代码3.3优化器参数切分配置 context.set_auto_parallel_context(parallel_mode=ParallelMode.AUTO_PARALLEL, enable_parallel_optimizer=True, device_num=4) class DenseMatMulNet(nn.Cell): def __init__(self): super(DenseMutMulNet, self).__init__() self.matmul1 = ops.MatMul.shard(((4, 1), (1, 1))) self.matmul2 = ops.MatMul.shard(((1, 1), (1, 4))) def construct(self, x, w, v): y = self.matmul1(x, w) z = self.matmul2(y, v) return z 3.1.8异构图切分 流水线切分、并行子图切分和算子级切分适用于模型的算子数量较大、同时参数较均匀地分布在各个算子中的情况。如果模型中的算子数量较少,同时参数只集中在几个算子中呢?Wide&Deep 就是这样的例子,如图3.19所示。Wide&Deep中的词表(Embedding table)作为需训练的参数可达几百GB甚至几TB,一方面,若放在加速器上执行,那么所需的加速器数量巨大,训练费用昂贵; 另一方面,若使用加速器计算,其获得的训练加速有限,同时会引发跨服务器的通信量,端到端的训练效率不会很高。 图3.19Wide&Deep模型的部分结构 仔细分析Wide&Deep模型的特殊结构后可知: 词表虽然参数量巨大,但其参与的计算量很少,可以将词表和其对应的算子词表查询(EmbeddingLookup)算子放置在Host端,利用CPU进行计算; 其余算子放置在加速器端。这样做能够同时发挥Host端内存量大、加速器端计算快的特性,同时利用了同一台服务器的Host到加速器高带宽的特性。图3.20展示了Wide&Deep异构切分的方式。 图3.20Wide&Deep异构切分方式 MindSpore支持异构图切分,如代码3.4所示,其中EmbeddingLookup算子配置属性为add_prim_attr('primitive_target','CPU'),即指示算子的执行位置是在Host端的CPU,其余算子的执行位置均是默认值(加速器)。MindSpore框架会自动将EmbeddingLookup的执行结果下沉到Host端; 类似地,在反向阶段计算EmbeddingLookup的梯度时,中间结果会由加速器端返还到Host端。 代码3.4MindSpore中异构切分的示例 context.set_auto_parallel_context(parallel_mode="semi_auto_parallel", device_num=4) class WideDeepSubNet(nn.Cell): def __init__(self, device_num): super(WideDeepSubNet, self).__init__() self.embedding_table = Parameter([self.vocab_size, self.embedding_size]) self.embedding_lookup = P.EmbeddingLookup().add_prim_attr('primitive_target', 'CPU') self.reduce_sum = P.ReduceSum().shard(((1, device_num, 1),)) def construct(self, in_hldr): wide_id = self. embedding_lookup (self.embedding_table, in_hldr) z = self. reduce_sum(wide_id) return z 3.1.9重计算 MindSpore根据正向图计算流程来自动推导出反向图,正向图和反向图一起构成了完整的计算图。在计算某些反向算子时,可能需要用到某些正向算子的计算结果,导致这些正向算子的计算结果需要驻留在内存中直到这些反向算子计算完,它们所占的内存才会被其他算子复用。这些正向算子的计算结果长时间驻留在内存中,会推高计算的内存占用峰值,这在大规模网络模型中尤为显著。 为了降低内存峰值,重计算技术可以不保存正向激活层的计算结果,让该内存可以被复用,然后在计算反向部分时,重新计算出正向激活层的结果。MindSpore提供了重计算的能力。 重计算功能具体实现方法是: 根据用户指定的需要做重计算的正向算子,复制出一份相同的算子,输出到反向算子上,并删除原正向算子与反向算子间的连边关系。另外,需要保证复制出来的算子,在计算相应的反向部分时才开始被计算,所以需要插入控制依赖去保证算子执行顺序。开启重计算功能前后的正反向示意图如图 3.21所示。 图3.21开启重计算功能前后的正反向示意图 为了方便用户使用,MindSpore目前不仅提供了针对单个算子设置的重计算接口,还提供了针对子模型设置的重计算接口。当用户调用子模型的重计算接口时,这个子模型中的所有正向算子都会被设置为重计算。以LeNet模型为例,重计算配置如代码3.5所示。 代码3.5LeNet模型重计算配置 import mindspore.nn as nn from mindspore.ops import operations as P class LeNet(nn.Cell): def __init__(self): super(LeNet, self).__init__() self.relu = P.ReLU() self.relu.recompute(True) self.batch_size = 32 self.conv1 = nn.Conv2d(1, 6, kernel_size=5, stride=1, padding=0, has_bias=False, pad_mode='valid') self.conv2 = nn.Conv2d(6, 16, kernel_size=5, stride=1, padding=0, has_bias=False, pad_mode='valid') self.pool = nn.MaxPool2d(kernel_size=2, stride=2) self.pool.recompute(True) self.reshape = P.Reshape() self.fc = nn.Dense(400, 120) 以GPT3模型为例,设置策略为对每层网络对应的子模型设置为重计算,然后每层网络的输出算子设置为非重计算。开启重计算功能前后GPT3内存使用比较如图3.22所示。 图3.22开启重计算功能前后的GPT3内存使用比较 第一个内存占用峰值为计算反向部分产生的,开启重计算功能后,该峰值被削平了,这得益于重计算带来的高内存复用率。 3.1.10GPT3超大规模分布式并行方案 GPT3是由OpenAI在2020年提出的语言模型,其由Transformer中的Encoder堆叠而成,参数量达1750亿个。训练此模型需要成百甚至上千台加速器,如何将这些加速器组织好,并且如何选择合理的并行方式使得模型训练能够充分利用带宽是难题。 常用的拓扑是两层架构: 服务器通过机架交换机相连,机架与机架之间通过核心交换机相连。网络带宽由大到小分别是: 同一服务器的不同加速器之间、同一机架的不同服务器上的加速器之间、不同机架的加速器之间。为了能够充分利用带宽,MindSpore采用了图3.23的并行方式: 不同机架之间是数据并行(每个机架都有完整的模型),机架内部的服务器之间是流水线并行(图3.23将模型划分为3个阶段),同一服务器的不同加速之间是算子级并行(即将张量切为4份)。这样做的原因是,算子级并行所需的通信量最大,流水线并行次之,数据并行所需的通信量最小。 图3.23GPT3并行策略配置 相比于传统训练GPT3的方式,此并行方式有两点不同: ①如前所述,图切分是在图编译阶段进行的,这样做有利于后续算子融合等优化点的实施; ②使用如图3.8所示的流水线并行方式,使得内存的峰值降低,可以跑下更大的批数据。 3.1.11小结 相比于其他开源框架,原生MindSpore在并行训练方面提供了较丰富的技术点。当用户遇到大模型训练的问题时,使用MindSpore能够避免自己陷入并行处理性能的麻烦,MindSpore将并行的性能问题与模型的算法逻辑解耦开,帮助用户更聚焦于算法逻辑。 当前,MindSpore提供了基础的并行技术点,如何选择最优的并行方式仍然需要用户干预,例如流水线并行的切分、重计算配置、通信算子融合等。如图3.5所示,现有的代价模型和策略搜索只支持算子并行。如何通过框架提供的全自动并行配置方案将用户彻底解放出来,是一个巨大的挑战。未来,技术人员将朝着这个方向继续完善MindSpore自动并行的能力。 3.2二阶优化 深度学习训练过程可以看成损失函数损失值下降过程,合适的优化器可以让深度学习训练时间大大缩短。优化器可以分为一阶优化器和二阶优化器,目前业界主流使用的仍然是一阶优化器,二阶优化器因为单步训练时间过久而没有被广泛应用。近年来,将二阶优化应用到深度学习训练中有了理论突破,并取得了不错的结果。本节内容分为四部分: 优化器的背景介绍,MindSpore自研二阶优化器THOR的介绍,THOR的实践应用和总结。 3.2.1优化器背景介绍 假设训练样本数据集 D=x1,y1,…,xi,yi,…,xN,yN,xi∈X,yi∈Y 参数θ表述的深度神经网络模型为 y^=fx; θ,x∈X,y^∈Y 定义在模型输出和真实标签y之间的损失函数为 L(y,fx;θ) 网络参数学习的过程是最小化损失函数的过程 minθ L(y,fx;θ) 图3.24深度学习训练过程模拟 给定数据集、模型、损失函数后,深度学习训练问题归结为优化问题,深度神经网络训练优化问题参数规模巨大,需要大量的计算,难以计算出解析解。因此,该过程也常常被比喻成下山,如图3.24所示,一个人站在山顶的时候如何在有限视距内寻找最快路径下山呢? 优化器就是在做这件事情,业界的优化算法可分为一阶优化算法和二阶优化算法。下面简单介绍业界的优化器情况。 1. 一阶优化器 梯度下降算法(Gradient Descent,GD)是机器学习中最经典的一阶优化算法,也是众多机器学习算法中最常用的优化算法。常用的一阶优化算法(如SGD算法)中对参数的更新采用以下规则: θ=θ-ηLθ 式中,θ为需要更新的参数,η为学习率,Lθ为损失函数对于参数的梯度。 但是主流随机梯度下降方法有以下问题: 太小的学习率会导致网络收敛过于缓慢; 学习率太高可能会影响收敛,并导致损失函数在最小值上波动,甚至出现发散,对参数比较敏感,容易收敛到局部最优,难以跳出鞍点。 因此,业界提出了很多随机梯度下降方法的改良算法,例如Momentum、Nesterov、AdaGrad、RMSprop、Adadelta和Adam等。这些改进后的优化算法可以利用随机梯度的历史信息来自适应地更新步长,使得它们更容易调参,而且方便使用。下面简单介绍业界常用的一阶优化算法Momentum和Adam。 1) Momentum Momentum是一阶优化算法,它对SGD进行了改进,在梯度下降的过程中加入了惯性,使得梯度在方向不变的维度上速度变快,在方向有所改变的维度上的更新速度变慢,与SGD相比可以加快收敛并减少震荡。其更新规则如下: vt=γvt-1+ηtLθ (3.4) θt=θt-1-vt (3.5) 式中,γ为动量超参;0≤γ<1; ηt为在第t个step时的学习率; θ是需要更新的参数; Lθ是损失函数对于参数的梯度。 2) Adam Adam也是一种自适应学习率的算法,该算法使用了动量变量vt和一阶梯度的指数加权移动平均变量st,并在初始时将它们置为0,其更新规则如下: vt=β1vt-1+(1-β1)Lθ (3.6) st=β2st-1+(1-β2)(Lθ)2 (3.7) v^t=vt1-βt1,s^t=st1-βt2 (3.8) θt=θt-1-ηtv^ts^t+ε (3.9) 式中,β1和β2为超参,通常取为0.9和0.999,ηt为在第t个step时的学习率,θ是需要更新的参数,Lθ是损失函数对于参数的梯度,ε是一个很小的数,为了避免分母为0,一般为1×10-8。 2. 二阶优化器 二阶优化算法利用目标函数的二阶导数进行曲率校正来加速一阶梯度下降。与一阶优化算法相比,其收敛速度更快,能高度逼近最优值,几何上下降路径也更符合真实的最优下降路径。 图3.25不同优化器下降路径 例如,二阶优化算法中的牛顿法就是用一个二次曲面去拟合当前所处位置的局部曲面,而梯度下降法是用一个平面去拟合当前的局部曲面。通常情况下,二次曲面的拟合会比平面更好,所以牛顿法选择的下降路径会更符合真实的最优下降路径。如图3.25所示,左边下降路径表示牛顿法的下降曲线,右边表示一阶梯度的下降曲线,二阶算法与一阶算法相比,可以更快地走到目的地,从而加速收敛。 从数学公式上来看,与一阶优化算法相比,二阶优化算法则是先将θ与一个矩阵G-1相乘,产生如下的更新规则 θ=θ-ηG-1Lθ 式中,G为二阶信息矩阵,不同的二阶优化算法中的G定义是不同的,常见的二阶优化算法有牛顿法、自然梯度法,分别对应的二阶信息矩阵G为海森矩阵、费雪矩阵。 牛顿法有着很好的局部收敛性质,当函数L在最优值点θ*点满足Lθ*=0,2Lθ*是正定矩阵且海森矩阵在极值点附近是李普希兹连续时,牛顿法二次收敛到最优值点。海森矩阵是一个由多变量实值函数的所有二阶偏导数组成的方块矩阵。海森矩阵可以表示为 Hij=2Lθiθj 式中,L即为损失函数,θ是需要更新的参数。 在SGD中,参数空间和函数空间的度量用的都是欧几里得距离,但欧几里得距离在一些情况下不能作为函数空间准确的距离度量。例如,在神经网络中,参数引起的目标函数变化是概率的变化,这并不适合在欧几里得空间度量,它不是概率属性变化的合理表征。KL散度是分布之间距离的合理度量。当使用KL散度作为概率分布之间距离的度量时,参数更新用到的梯度就是自然梯度。自然梯度法中的费雪矩阵可以表示为 F=ElogP(y|x,θ)θlogp(y|x,θ)Tθ 式中,P(y|x,θ)是网络模型的预测分布,p(y|x,θ)是其概率密度,θ是网络模型的参数。 二阶优化算法虽然收敛速度快,但是计算二阶矩阵的逆的时间复杂度为O(n3),当模型参数量为 nθ 时,对应的二阶信息矩阵的大小为 nθ×nθ。在深度学习模型中,nθ 常常在数百万的量级,此时二阶信息矩阵的逆无法计算。因此,如何降低二阶信息矩阵求逆的计算复杂度成为关键问题。接下来介绍下深度学习中的二阶优化器。 1) KFAC 克罗内克因子近似曲率(Kroneckerfactored Approximate Curvature,KFAC)是由Martens等学者于2015年提出的基于自然梯度法的二阶优化算法,该算法高效地近似了自然梯度法中的费雪矩阵,并且给出了充分的理论证明,给后续许多二阶优化算法带来了启发。这里重点介绍是如何做近似的,具体的理论推导可参考相关论文。 第一步,KFAC假设不同层之间的参数是独立的,从而将二阶信息矩阵按网络层解耦,近似成块对角矩阵,然后对二阶信息矩阵求逆就变为了对这些块对角矩阵求逆。 第二步,将按层解耦后的每个块矩阵又近似为两个小得多的矩阵的克罗内克乘积。克罗内克乘积公式如下: b11b12 b21b22c11c12c13 c21c22c23 c31c32c33=b11c11b11c12b11c13b12c11b12c12b12c13 b11c21b11c22b11c23b12c21b12c22b12c23 b11c31b11c32b11c33b12c31b12c32b12c33 b21c11b21c12b21c13b22c11b22c12b22c13 b21c21b21c22b21c23b22c21b22c22b22c23 b21c31b21c32b21c33b22c31b22c32b22c33 (3.10) 图3.26KFAC对二阶矩阵的近似 由于两个矩阵的克罗内克乘积的逆等于它们逆的克罗内克乘积,且这两个较小矩阵会比整个块矩阵更容易计算和求逆,通过这样的近似和分解,KFAC大大简化了费雪矩阵的计算。 比如,对于一个342的神经网络,其结构示意图如图3.26所示。 原始二阶信息矩阵维度为20×20,经过第一步近似变为了一个12×12和8×8的矩阵,再经过第二步近似,分别变为了3×3和4×4、4×4和2×2维度的矩阵,如图3.27所示。 图3.27KFAC对二阶矩阵的近似 图3.28KFAC实验结果图 在CIFAR10上对该算法进行实验,使用cudaconvnet网络,该网络包含3层卷积层和1层全连接层,图3.28中横坐标表示运行时间,纵坐标表示分类误差,收敛速度较慢的曲线为SGD,收敛速度较快的曲线为KFAC(图3.28中的KFCpre即为KFAC),实线是测试误差,虚线是训练误差。KFAC达到19%的测试误差仅需3min,SGD需要9min。在训练误差上KFAC的优势更加明显,KFAC达到6%仅需4min,SGD需要30min。 KFAC使自然梯度法应用到了深度学习训练中,并取得了很好的结果,但是KFAC未考虑大型神经网络计算复杂度,例如对于大型神经网络ResNet50,全连接层维度为2048×1000,按照KFAC算法分解后,仍然要计算一个2048维和一个1000维的逆矩阵,求逆时间虽然大大降低,但是依然很慢。在ResNet50模型、ImageNet全集上,单次迭代所需时间是秒级别,单迭代所需时间是一阶的33倍。 2) SPNGD SPNGD(Scalable and Practical Natural Gradient Descent,可扩展且实用的自然梯度下降算法)是由Osawa等学者于2020年提出的基于自然梯度法的二阶优化器,并在大数据集ImageNet上训练ResNet50,使用1024块V100,仅需5.5min即可收敛到75.4%。 SPNGD受KFAC启发,在KFAC基础上主要做了以下两点改进。 第一点,SPNGD使用了过时的二阶信息矩阵,在训练过程中不是每个迭代都更新二阶信息矩阵的逆,设当前更新的二阶信息矩阵为X,上一次更新的二阶矩阵为X-1,上上次更新的二阶矩阵为X-2,过程为判断X是否与X-1相似,若不相似,减少下一次更新的间隔,若相似则进一步判断X是否与X-2相似,若不相似,更新间隔保持不变,若相似则增大下次更新间隔。具体计算规则如图3.29所示,算法1体现了如何使用过时二阶矩阵逆来更新梯度,通过控制更新间隔来变化地更新二阶矩阵逆; 算法2体现了更新间隔的变化方式。 图3.29SPNGD 更新规则 第二点,为了更高效地进行分布式训练,提出了数据和模型混合并行,具体做法如图3.30所示。 阶段1: 给每块GPU分配不同的数据集,做数据并行,每块GPU根据分配的数据集计算前向过程和前向过程中二阶信息矩阵需要的信息A。 阶段2: 对A进行归并,归并时使用ReduceScatterV,而不是常规使用的AllReduce,这个操作是指每个GPU只归并指定层中的A,将数据并行转为模型并行,并且在A通信时同时计算反向过程,进一步减少计算时间。 阶段3: 对反向过程中计算得到的相关信息也使用ReduceScatterV做归并。这时每块GPU有指定层的完整信息。 阶段4: 每块GPU分别计算指定层的二阶信息矩阵的逆,并更新网络参数。 阶段5: 对网络参数进行AllGatherV,同步所有GPU中的参数信息。 图3.30SPNGD混合并行模式 该算法还有其他改进点,如对BN层的二阶梯度的特殊处理,如何做数据增强,如何更好地调参,等等。图3.31展示了该算法的实验结果,Hardware表示硬件平台,Software是指使用的深度学习框架,Batch size是每次训练的图片数量,Optimizer表示使用的优化器,Steps是总共训练的步数,Time/step是每步训练时间,Time是总体训练时间,Accuracy是最后收敛精度。在ImageNet上训练ResNet50,使用1024块V100,仅需5.5min即可收敛到75.4%,该算法使二阶优化器在大数据集和大型网络中与一阶优化器比较取得了有竞争力的结果。 图3.31SPNGD实验结果 3.2.2THOR简介 当前业界已有的二阶优化算法计算量依然过大,与一阶相比优势不明显,且使用场景较为单一。MindSpore提出了自研算法THOR(Tracebased Hardwaredriven layerORiented Natural Gradient Descent Computation),该算法单步时间短,且在BERT和ResNet50中端到端时间优势均非常大。THOR主要做了以下两点创新。 1. 降低二阶信息矩阵更新频率 通过实验观察费雪矩阵的F范数(Frobenius norm)可知,其在前期变化剧烈,后期逐渐变稳定,从而假设 {Fk}nk=1是一个马尔可夫过程,可以收敛到一个稳态分布π,其中Fk代表第k个迭代时的费雪矩阵。因此,在训练过程中逐步增大费雪矩阵的更新间隔,可以在不影响收敛速度的情况下,减少训练时间。例如,在ResNet50中,更新间隔步数随着训练的进行越来越大,到后期每个周期只需更新一次二阶信息矩阵。 THOR受KFAC启发,将费雪矩阵按层解耦来降低矩阵复杂度,分别针对每一层的费雪矩阵做实验,发现有些层的费雪矩阵趋于稳态的速度更快,因此在统一的更新间隔上,更加细粒度地去调整每一层的更新频率。THOR使用矩阵的迹作为判断条件,当迹的变化情况大于某一阈值时更新该层的二阶信息矩阵,否则沿用上一个迭代的二阶信息矩阵,并且引入了停止更新机制,当迹的变化量小于某个阈值时停止更新该层二阶信息矩阵,具体更新公式如下: Δk=‖tr(Fki+λI)|-|tr(Fk-1i+λI)‖|tr(Fk-1i+λI)| (3.11) 更新Fki,Δk∈(w1,+∞) 不更新Fki,沿用上一个迭代的Fk-1i,Δk∈[w2,w1] 停止更新Fki,后期都使用Fk-1i,Δk∈[0,w2] (3.12) 2. 硬件感知矩阵切分 THOR在将费雪矩阵按层解耦的基础上,进一步假设每个网络层中的输入和输出块之间也是独立的,例如将每层网络的输入输出切分为n个块,这n个块之间即是独立的,根据该假设对二阶信息矩阵做进一步的切分,从而提高了计算效率。THOR结合矩阵信息损失数据和矩阵性能数据确定了矩阵分块维度,从而大大提升费雪矩阵求逆时间。 那么如何确定矩阵分块维度呢?具体方法如下: (1) 根据费雪矩阵中维度最大的那一层,确定矩阵切分维度,以ReseNet50为例,网络层中的最大维度为2048,确定矩阵切分维度为[1,16,32,64,128,256,512,1024,2048]。 (2) 根据确定的矩阵维度,借助谱范数计算每个维度下的矩阵损失,具体公式为 L=1-λmax(A^A^T)λmax(AAT) (3.13) 式中,λmax(X)表示矩阵X的最大特征值,A表示原始未分割矩阵,A^表示分割后的矩阵。然后统计在该维度下损失小于1%的矩阵数量,最后通过除以总的矩阵数量得到标准化后的矩阵损失信息。 (3) 根据确定的矩阵维度,计算每个维度下的矩阵求逆时间,再通过公式normalizedn=p1pn得到每个维度下标准化后性能数据,其中p1表示维度最小的矩阵的性能数据,pn表示第n个维度下的性能数据。 图3.32切分维度确定示意图 (4) 根据标准化后的矩阵损失信息和标准化后的性能数据绘图,如以ResNet50为例,可得到图3.32。图3.32中下降曲线为性能曲线,上升曲线表示矩阵损失曲线,图3.32中曲线交叉点为106,与128最接近,最后确定矩阵切分维度为128。 3. 实验结果 图3.33展示了THOR在ResNet50+ImageNet、Batch size为256时一二阶上的训练曲线图,其中train loss表示训练误差,test accuracy表示测试精度,epoch表示迭代数,wallclock time表示时间,其中下降较快的曲线和上升较快的曲线是本算法曲线,另外差距较明显的曲线是Momentum的训练曲线。 图3.33THOR在ResNet50上的实验结果 图3.33中的THOR、THOR_stop、THOR_NT分别表示(w1,w2)=(0.01,0),(w1,w2)=(0.01,0.001),(w1,w2)=(0,0),从图3.33中可以看到THOR收敛所需迭代数大约是一阶的一半,且单step的时间与一阶相差也不大。相比一阶算法需要117min,二阶优化器端到端时间提速约40%。 THOR还测试了在不同Batch size下ResNet50+ImageNet的收敛结果,结果见图3.34,其中Hardware表示硬件平台,Software是指使用的深度学习框架,Batch size是每次训练的图片数量,Optimizer表示使用的优化器,Time指总体训练时间,Accuracy是指最后收敛精度。当Batch size为8192、使用256块Ascend 910时,只需2.7min精度即可收敛到75.9%。 图3.34不同Batch size下THOR在ResNet50上的实验结果 在BERT+WIkipedia中,THOR也有不错的表现效果,以MLPerf为标准,精度达到71.2%,与一阶相比端到端提升30%,实验结果见图3.35。图3.35中横坐标表示训练时间,纵坐标表示测试精度,上升较快的曲线是THOR的训练曲线,另一条为Lamb的训练曲线。 图3.35THOR在BERT上的实验结果 3.2.3THOR的实践应用 THOR算法的部分内容当前已经在MindSpore中开源,源码位置: https://link.zhihu.com/?target=https%3A//gitee.com/mindspore/mindspore/blob/master/mindspore/nn/optim/thor.py。 MindSpore中使用THOR训练网络非常简单,只需要四行代码即可调用THOR训练网络。THOR训练网络代码如代码3.6所示。 代码3.6THOR训练网络代码 from mindspore.nn.optim import thor as THOR #引用二阶优化器 #创建网络 net = Net() #调用优化器 opt = THOR(net, lr, Tensor(damping), config.momentum, config.weight_decay, config.loss_scale, config.batch_size, split_indices=split_indices, frequency=config.frequency ) #增加计算图提升性能 model = ConvertModelUtils().convert_to_thor_model(model=model, network=net, loss_fn=loss, optimizer=opt, oss_scale_manager=loss_scale, metrics={'acc'}, amp_level="O2", keep_batchnorm_fp32=False) #训练网络 model.train(config.epoch_size, dataset, callbacks=cb, sink_size=dataset.get_dataset_size(), dataset_sink_mode=True) (1) 导入MindSpore所需的二阶优化器的包,位于 mindspore.nn.optim。 (2) 创建网络和定义THOR优化器,传入网络信息和THOR所需的超参信息(如学习率、正则化项系数等)。 (3) 调用convert_to_thor_model函数,该函数通过增加计算图使THOR达到更优性能。 (4) 调用model.train开始训练网络。 1. 源码分析 __init__函数用于THOR的初始化,需要传入THOR所需的超参和网络结构,THOR支持GPU和Ascend,分别为class THOR_GPU(Optimizer)和class THOR_Ascend(Optimizer),这两个类之间的主要差别是算子不同。下面以class THOR_Ascend(Optimizer)为例进行分析,如代码3.7所示。 代码3.7初始化THOR源码 class THOR_Ascend(Optimizer): def __init__(self, net, learning_rate, damping, momentum, weight_decay=0.0, loss_scale=1.0, batch_size=32, decay_filter=lambda x: x.name not in [], split_indices=None): params = filter(lambda x: x.requires_grad, net.get_parameters()) #THOR初始化,定义算子和变量 super(THOR_Ascend, self).__init__(learning_rate, params, weight_decay, loss_scale) if isinstance(momentum, float) and momentum < 0.0: raise ValueError("momentum should be at least 0.0, but got momentum {}".format(momentum)) self.momentum = Parameter(Tensor(momentum, mstype.float32), name="momentum") self.params = self.parameters self.moments = self.params.clone(prefix="moments", init='zeros') self.hyper_map = C.HyperMap() self.opt = P.ApplyMomentum() self.net = net self.matrix_A_cov = ParameterTuple(filter(lambda x: 'matrix_A' in x.name, net.get_parameters())) self.matrix_G_cov = ParameterTuple(filter(lambda x: 'matrix_G' in x.name, net.get_parameters())) ... MindSpore中所有优化器都继承了class Optimizer,该基类中定义了一些基本函数(如获取学习率、梯度缩放等)。THOR初始化时将传进去的超参定义为类属性方便调用,并且定义了后续计算会使用到的算子。 也就是说,初始化函数的作用就是定义THOR计算所需要用到的算子和变量(Parameter、Tensor等)。 下面重点介绍self.matrix_A_cov和self.matrix_G_cov。这两个变量是计算二阶梯度所需要的信息,分别为每层输入a的协方差矩阵A=aTa和每层输出的一阶导数g的协方差矩阵G=gTg,其中a、g已经在运行时的前向过程和反向过程中保存下来,该源码地址为https://gitee.com/mindspore/mindspore/blob/master/mindspore/nn/layer/thor_layer.py。 创建THOR时的入参: (1) net: 本次训练建立的模型; (2) learning_rate: 学习率超参; (3) damping: 二阶矩阵中加的正则化项的超参; (4) momentum: 动量超参; (5) weight_decay: 权值衰减,用于防止过拟合,默认值为0.0,即不使用权值衰减; (6) loss_scale: 用于缩放训练过程中的loss,防止梯度越界,默认值为1.0,即不使用缩放; (7) batch_size: 当前训练一个step所使用的数据量,默认为32; (8) decay_filter: 选择对哪些层做weight decay,当weight_decay>0时起作用; (9) split_indices: 这个参数用于加速allreduce过程。 _get_Ainv_Ginv_Amax_Gmax_list 函数用于计算协方差矩阵A/G的逆,并返回求逆后的矩阵。具体过程是遍历模型所有层,按层处理,对每一层的协方差矩阵加上正则化项,然后对矩阵进行Cholesky分解从而来求逆。当前开源代码THOR中支持全连接层和卷积层的处理。二阶矩阵求逆源码如代码3.8所示。 代码3.8二阶矩阵求逆源码 def _get_Ainv_Ginv_Amax_Gmax_list(self, gradients, damping_step, matrix_a_allreduce, matrix_g_allreduce, matrix_a_max_allreduce, matrix_g_max_allreduce): """get matrixA inverse list, matrixG inverse list, matrixA_max list, matrixG_max list""" #计算A/G的逆 for i in range(len(self.params)): thor_layer_count = self.weight_fim_idx_map[i] conv_layer_count = self.weight_conv_idx_map[i] layer_type = self.weight_layerType_idx_map[i] if layer_type in [Conv, FC, Embedding]: g = gradients[i] matrix_A = self.matrix_A_cov[thor_layer_count] matrix_G = self.matrix_G_cov[thor_layer_count] matrix_A = F.depend(matrix_A, g) matrix_G = F.depend(matrix_G, g) A_shape = self.shape(matrix_A) A_eye = self.eye(A_shape[0], A_shape[0], mstype.float32) G_shape = self.shape(matrix_G) G_eye = self.eye(G_shape[0], G_shape[0], mstype.float32) if layer_type == Conv: ... elif layer_type == FC: matrix_A = matrix_A + damping * A_eye matrix_A_inv = self.cholesky(matrix_A) #求逆 matrix_A_inv = self.vector_matmul(matrix_A_inv, matrix_A_inv) _get_second_gradients函数用于计算最终参数更新方向,参数更新方向公式为(Aki-1+λI)-1(Gki+λI)-1vec(WiJk),而(AG)vec(X)=vec(GXAT),所以代码实际实现的公式为(Gki+λI)-1WiJk(Aki-1+λI)-1,其代码如代码3.9所示。 代码3.9二阶梯度计算源码 def _get_second_gradients(self, new_grads, damping_step, gradients): """get second gradients for thor""" #计算参数更新方向 params_len = len(self.params) for i in range(params_len): ... else: ... elif layer_type == FC: temp_a = self.matrix_A_cov[thor_layer_count] temp_g = self.matrix_G_cov[thor_layer_count] temp_a = self.cast(temp_a, mstype.float16) temp_g = self.cast(temp_g, mstype.float16) g = self.cast(g, mstype.float16) g = self.matmul(temp_g, g) #一阶梯度左乘G g = self.matmul(g, temp_a) #右乘A,得到最终更新方向 g = self.cast(g, mstype.float32) construct函数是在网络训练过程中会实际执行的内容,该函数中包含了上述两个函数_get_Ainv_Ginv_Amax_Gmax_list和_get_second_gradients的调用,该函数完成了二阶矩阵的计算和梯度更新方向的调整,其代码如代码3.10所示。 代码3.10THOR计算和梯度更新方向源码 def construct(self, gradients): #二阶矩阵计算和梯度更新方向 params = self.params moments = self.moments damping_step = self.gather(self.damping, self.cov_step, self.axis) damping_step = self.cast(damping_step, mstype.float32) if self.thor: matrix_A_allreduce = () matrix_G_allreduce = () matrix_A_max_allreduce = () matrix_G_max_allreduce = () matrix_A_allreduce, matrix_G_allreduce, matrix_A_max_allreduce, matrix_G_max_allreduce = self._get_Ainv_Ginv_Amax_Gmax_list(gradients, damping_step, matrix_A_allreduce, matrix_G_allreduce, matrix_A_max_allreduce, matrix_G_max_allreduce) #计算A/G的逆 ... new_grads = () for i in range(len(self.params)): ... if self.conv_layer_count > 0:#有卷积层时的处理 ... else: #都是全连接层时的处理 if layer_type == Embedding: ... elif layer_type == FC: temp_a = matrix_A_allreduce[thor_layer_count] temp_g = matrix_G_allreduce[thor_layer_count] fake_A = self.assign(self.matrix_A_cov[thor_layer_count], temp_a) fake_G = self.assign(self.matrix_G_cov[thor_layer_count], temp_g) g = F.depend(g, fake_A)#确保执行顺序 g = F.depend(g, fake_G) temp_a = self.cast(temp_a, mstype.float16) temp_g = self.cast(temp_g, mstype.float16) g = self.cast(g, mstype.float16) g = self.matmul(temp_g, g) g = self.matmul(g, temp_a)#将一阶方向变为二阶方向 g = self.cast(g, mstype.float32) elif layer_type == LayerNorm: g = self._process_layernorm(damping_step, g) new_grads = new_grads + (g,) gradients = new_grads #计算后得到的更新方向 else: #该分支表示使用过时二阶信息更新参数 new_grads = () gradients = self._get_second_gradients(new_grads, damping_step, gradients) #调用_get_second_gradients函数计算方向 … 2. THOR的实践应用 THOR在ResNet50和BERT上的脚本已开源,链接如下: ResNet50: https://gitee.com/mindspore/mindspore/blob/master/model_zoo/official/cv/resnet/train.py BERT: https://gitee.com/mindspore/mindspore/blob/master/model_zoo/official/nlp/bert/run_pretrain.py 1) ResNet50 优化器的调用方式与前面提到的一致,在这个例子中展开具体训练过程。 (1) 创建网络训练需要的训练集和网络定义为ResNet50。 (2) 设置THOR需要用到的超参策略,其他超参值设定可在该目录下的src/config.py中修改。 (3) 创建THOR优化器,并传入设置的超参值。 (4) 转换模型保存二阶所需信息。 (5) 训练网络。 代码3.11THOR在ResNet上的示例代码 from mindspore.nn.optim import thor as THOR #引用二阶优化器 from src.resnet import resnet50 as resnet from mindspore.train.model import Model ... if __name__ == '__main__': ... #创建网络训练过程中的训练集 dataset = create_dataset(dataset_path=args_opt.dataset_path, do_train=True, repeat_num=1, batch_size=config.batch_size, target=target, distribute=args_opt.run_distribute) step_size = dataset.get_dataset_size() #创建ResNet50模型 net = resnet(class_num=config.class_num) ... # init lr if cfg.optimizer == "Thor": #设置超参值 from src.lr_generator import get_thor_lr lr = get_thor_lr(0, config.lr_init, config.lr_decay, config.lr_end_epoch, step_size, decay_epochs=39) # define loss, model if target == "Ascend": if args_opt.dataset == "imagenet2012": if not config.use_label_smooth: config.label_smooth_factor = 0.0 loss = CrossEntropySmooth(sparse=True, reduction="mean", smooth_factor=config.label_smooth_factor, num_classes=config.class_num) else: loss = SoftmaxCrossEntropyWithLogits(sparse=True, reduction='mean') loss_scale = FixedLossScaleManager(config.loss_scale, drop_overflow_update=False) #高层抽象,集成网络模型的训练和测试 model = Model(net, loss_fn=loss, optimizer=opt, loss_scale_manager=loss_scale, metrics={'acc'}, amp_level="O2", keep_batchnorm_fp32=False) if cfg.optimizer == "Thor" and args_opt.dataset == "imagenet2012": from src.lr_generator import get_thor_damping #设置超参damping damping = get_thor_damping(0, config.damping_init, config.damping_decay, 70, step_size) #用于通信时的并行加速 split_indices = [26, 53] #创建THOR优化器 opt = THOR(net, lr, Tensor(damping), config.momentum, config.weight_decay, config.loss_scale, config.batch_size, split_indices=split_indices, frequency=config.frequency) #增加计算图提升性能 model = ConvertModelUtils().convert_to_thor_model(model=model, network=net, loss_fn=loss, optimizer=opt, loss_scale_manager=loss_scale, metrics={'acc'}, amp_level="O2", keep_batchnorm_fp32=False) ... #训练网络 model.train(config.epoch_size - config.pretrain_epoch_size, dataset, callbacks=cb, sink_size=dataset.get_dataset_size(), dataset_sink_mode=dataset_sink_mode) 输入 bash run_distribute_train.sh resnet50 imagenet2012 [RANK_TABLE_FILE] [DATASET_PATH] [PRETRAINED_CKPT_PATH] 即可运行脚本。 2) BERT BERT中的步骤与ResNet50类似。 (1) 创建网络训练需要的训练集和网络定义为BERT。 (2) 设置THOR需要用到的超参策略,其他超参值设定可在该目录下的src/config.py中修改。 (3) 优化器创建时传入BERT设定的超参值,本例中创建时传入decay_filter=lambda x: 'layernorm' not in x.name.lower() and 'bias' not in x.name.lower() 表示做weight decay操作时排除LN层和FC中的bias参数。 (4) 转换模型保存二阶所需信息。 (5) 训练网络。 代码3.12THOR在BERT上的示例代码 from mindspore.nn.optim import thor as THOR #引用二阶优化器 from src import BertNetworkWithLoss ... def _get_optimizer(args_opt, network): """get bert optimizer, support Lamb, Momentum, AdamWeightDecay.""" if cfg.optimizer == 'Lamb': ... elif cfg.optimizer == "Thor": from src.utils import get_bert_thor_lr, get_bert_thor_damping #设置lr和damping的超参值 lr = get_bert_thor_lr(cfg.Thor.lr_max, cfg.Thor.lr_min, cfg.Thor.lr_power, cfg.Thor.lr_total_steps) damping = get_bert_thor_damping(cfg.Thor.damping_max, cfg.Thor.damping_min, cfg.Thor.damping_power, cfg.Thor.damping_total_steps) split_indices = None #设置并行加速方式 if bert_net_cfg.num_hidden_layers == 12: if bert_net_cfg.use_relative_positions: split_indices = [29, 58, 87, 116, 145, 174, 203, 217] else: split_indices = [28, 55, 82, 109, 136, 163, 190, 205] elif bert_net_cfg.num_hidden_layers == 24: if bert_net_cfg.use_relative_positions: split_indices = [30, 90, 150, 210, 270, 330, 390, 421] else: split_indices = [38, 93, 148, 203, 258, 313, 368, 397] #创建优化器 optimizer = THOR(network, lr, damping, cfg.Thor.momentum, cfg.Thor.weight_decay, cfg.Thor.loss_scale, cfg.batch_size, decay_filter=lambda x: 'layernorm' not in x.name.lower() and 'bias' not in x.name.lower(), split_indices=split_indices, frequency=cfg.Thor.frequency) ... return optimizer def run_pretrain(): ... #创建数据集 ds = create_bert_dataset(device_num, rank, args_opt.do_shuffle, args_opt.data_dir, args_opt.schema_dir) #创建网络和损失函数 net_with_loss = BertNetworkWithLoss(bert_net_cfg, True) ... #加载初始checkpoint if args_opt.load_checkpoint_path: param_dict = load_checkpoint(args_opt.load_checkpoint_path) load_param_into_net(net_with_loss, param_dict) #动态loss缩放 if args_opt.enable_lossscale == "true": ... #固定loss缩放值 else: #BERT网络训练封装 net_with_grads = BertTrainOneStepCell(net_with_loss, optimizer=optimizer) #创建网络 model = Model(net_with_grads) #增加计算图提升性能 model = ConvertModelUtils().convert_to_thor_model(model, network=net_with_grads, optimizer=optimizer) #网络训练 model.train(new_repeat_count, ds, callbacks=callback, dataset_sink_mode=(args_opt.enable_data_sink == "true"), sink_size=args_opt.data_sink_steps) if __name__ == '__main__': set_seed(0) run_pretrain() 输入bash scripts/run_distributed_pretrain_ascend.sh /path/cnwiki128 /path/hccl.json即可运行脚本。 3.2.4小结 二阶优化算法利用曲率信息可以更快收敛,但因为单步计算时间过久,因此没有被广泛应用。如何降低二阶算法的单步计算时间是当前业界非常关注的事情,近年来,二阶优化也取得了许多理论突破,如KFAC、SPNGD等算法,但是这些算法计算时间依然很长,与主流一阶相比优势不明显。MindSpore提出了自研优化器THOR,该算法代码已开源。THOR通过降低二阶信息矩阵更新频率和硬件感知矩阵切分,在不影响收敛的情况下,大大降低了计算量,在CV经典模型ResNet50上与主流一阶优化器相比端到端时间减少40%,在NLP经典模型BERT上减少30%。未来THOR还可应用到更大规模的模型上(如GPT3)和其他领域中(如强化学习)。 3.3模型量化 随着物联网的兴起和移动设备的普及,越来越多的应用选择在边缘设备和移动设备上使用深度学习技术。一方面,硬件资源的限制给深度神经网络在这些端侧设备上的推理性能和效率带来了不小的挑战; 另一方面,为了满足各种AI应用的要求,深度神经网络正在往结构更复杂、参数更多的方向发展。复杂的网络固然具有精度优势,但巨大的存储空间和计算资源消耗也使其难以在端侧设备上进行应用和部署。模型压缩就是应对该类问题衍生出的技术之一。模型压缩技术能够有效地降低网络冗余参数,从而减少模型存储占用、通信带宽和计算复杂度,有助于深度神经网络在端侧设备上的应用部署。 常见的模型压缩技术包括轻量化设计、网络结构融合、模型量化、模型剪枝。模型量化作为一种能够有效减少模型大小和加速推理的优化技术,已经得到了学术界和工业界的广泛研究和应用。早在1990年就有Fiesler等人对神经网络量化的研究,如今8bit整数的运算性能提升已经成为各大芯片厂商竞逐的重点之一,TensorFlow、TensorRT和MindSpore等训练和推理框架也已经支持模型量化的功能。 根据量化过程中是否需要进行训练来区分,一般可将模型量化分为感知量化训练和训练后量化两种类型。 (1) 感知量化训练(quantization aware training): 在模型训练过程中加入量化相关处理。感知量化训练需要训练数据,在模型准确率上通常表现更好,适用于对模型压缩率和量化精度要求较高的场景。 (2) 训练后量化(posttraining quantization): 在模型训练后进行量化,训练后量化简单易用,只需少量校准数据,适用于追求高易用性和缺乏训练资源的场景。 3.3.1量化算法原理 计算机中不同数据类型占用的比特位数和表示数据范围也不相同。如表3.1所示,比特位数越小,占用存储越小,同时能表示的数值范围也越小。深度神经网络的权重参数一般用单精度浮点数(float32)表示和存储。如果用有符号8位整数(int8)来近似表示浮点型参数,权重参数对应的存储大小就能减小到原来的四分之一。将浮点型的模型权重或激活值用更少比特位数的整数近似表示,通过牺牲一定的精度为代价,来达到减少模型大小和加快推理速度等目的的过程就是模型压缩中的量化。 表3.1不同数据类型的比特位数和数值范围 数据类型 比特位数 数值范围 单精度浮点数(float32) 32 -3.4×1038~3.4×1038 半精度浮点数(float16) 16 -65504~65504 有符号8位整数(int8) 8 -128~127 有符号4位整数(int4) 4 -8~7 一般小于或等于8bit的量化称为低比特量化。8bit量化是当前工业界最常见的量化方法。主流的深度学习框架也已支持8bit量化功能,许多测试数据均表明,将网络前向过程的浮点计算替换成8bit计算,不会带来明显的精度下降。1bit量化是模型压缩的极致。量化后取值为0/1或者1/-1,也称二值量化。二值量化在计算时可以使用高效的XNOR和BitCount位运算代替浮点运算。BNN和XNORNet是二值量化的开创性工作,XNORNet中验证了使用1bit卷积运算在CPU上的计算速度可达到卷积浮点运算的58倍。 根据量化的等级间距离,量化可分为均匀(uniform)量化和非均匀(nonuniform)量化。其中均匀量化的等级间距离相等,又称为线性量化; 非均匀量化的等级间距离不相等,又称为非线性量化,如对数量化或指数量化。深度神经网络的权重和激活值的数据范围大多是非均匀分布,即数据中存在部分区域密集和部分区域稀疏。理论上非均匀量化能达到更好的精度,但该方法的计算复杂度相对较高,不利于硬件加速。目前线性量化是在工业界落地的常用方法。下面主要介绍线性量化的算法原理。 假设r∈R表示浮点数,q∈Q表示量化后的整数(R和Q表示浮点数和整数的集合,r和q表示集合中的一个数)。浮点数r和整数q的转换公式如下: q=cliproundrs+z,qmin,qmax (3.14) r=s×(q-z) (3.15) 式(3.14)称为量化,式(3.15)则称为反量化。式中,s是量化因子,表示从浮点数到整数的量化步长,一般为float32类型; z是零点,表示浮点数0对应量化后的值,和q为相同数据类型; qmax和qmin表示量化后整数类型的最大值和最小值,例如量化为无符号8位整数(uint8)时,qmin和qmax取值分别为0和255; round为取整函数; clip(x,qmin,qmax)函数用于限制返回值在[qmin,qmax]范围内,如果x小于qmin,则返回qmin,若x大于qmax,则返回qmax。量化参数s和z的计算公式如下: s=rmax-rminqmax-qmin (3.16) z=roundqmax-rmaxs (3.17) 式中,rmin和rmax表示R中的最小值和最大值,它们的取值与量化后推理精度密切相关。如果R中存在少量数据是边界离群点,直接取R的最小值和最大值作为rmin和rmax的取值,按照上述公式进行量化会引起密集区域数据信息丢失从而导致量化误差。反之,如果缩小rmin和rmax的取值范围会引起边界区域数据信息丢失从而导致截断误差。如何选取合适的rmin和rmax以提高准确率也是业界的一个研究方向,目前主流的计算rmin和rmax的方法有直接计算最小值和最大值、KL散度校准、指数移动平均值统计、在训练过程中学习更新等。 当式(3.17)的z等于0时,又称为对称(symmetric)量化; 反之,z不等于0时称为非对称(asymmetric)量化。对称量化和非对称量化各有优势。相比而言,对称量化中z=0在公式上可简化相关计算步骤,降低推理时的计算复杂度。非对称量化对rmin/rmax的取值无特殊要求,能够更加充分利用量化比特位,量化损失相对会更少。在实际应用时,权重的分布一般是关于0对称的高斯分布,适合采用对称量化; 而ReLU和ReLU6后的激活值为非负值,适合采用非对称量化。 按照量化参数s和z的共享范围,又可以分为逐层(perlayer)量化和逐通道(perchannel)量化。逐层量化以一个网络层为单位,每层共同使用一组量化参数; 逐通道量化以网络层中的一个通道为单位,每层的每个通道单独使用一组量化参数。逐通道量化粒度更细,可以获得更高的量化精度,但计算量也相对更复杂。 根据深度神经网络中的量化对象来分,量化可以分成如下几类: (1) 权重(weight)量化: 对网络层的权重进行量化。由于模型训练后参数需要进行存储,对权重参数进行量化可以直接压缩模型大小。在实际应用中,对权重量化时一般采用对称量化,直接统计权重的rmin和rmax,按照量化精度要求选取逐层或者逐通道量化。 (2) 激活值(activation)量化: 对网络层的输出进行量化。由于网络层的输入也可以看成是上一层的输出,因此对于网络层的输入和输出量化统一称为激活值量化。激活值量化可以减小计算过程中的内存占用。如果进一步结合权重量化,则可以充分利用硬件平台的整数计算性能从而提升推理性能。在实际应用中,激活值量化一般采用非对称量化和逐层量化。由于激活值的rmin和rmax通常不容易提前获得,因此需要在训练或者推理过程中进行动态计算,相对熵散度校准和指数移动平均值是常用的两种方法。 (3) 梯度(gradient)量化: 对训练过程中的梯度进行量化,主要用于在分布式训练中提升训练速度。例如,联邦学习中通过梯度量化来提高计算通信比。 下面以矩阵乘法为例,介绍矩阵乘法积中权重和激活值量化计算过程。假设有Y=W×X,其中wi,j∈W、xj,k∈X、yi,k∈Y都是浮点数,则yi,k可表示为 yi,k=∑jwi,jxj,k (3.18) 如果对权重W和激活值X/Y进行逐层量化,sy/zy、sw/zw和sx/zx分别为Y、W和X的量化参数,可得 sy(yi,kq-zy)=∑jsw(wi,jq-zw)sx(xj,kq-zx) (3.19) 式中,带下标q的值表示对应量化后的整数值。进一步化简后可得 yi,kq=zy+swsxsy∑jwi,jq-zw(xj,kq-zx) (3.20) 式(3.20)即为矩阵乘法积量化公式。如果采用对称量化,可以进一步简化为 yi,kq=swsxsy∑jwi,jqxj,kq (3.21) 在式(3.20)和式(3.21)中,除swsxsy外都是整数运算。量化后的权重用整数wi,jq表示,结合硬件加速器的整型运算性能优势,可以达到减小内存占用和提升推理性能的目的。 综上所述,模型量化能带来很多好处,但与此同时也面临着巨大挑战,其中最核心的挑战是如何在尽量减小模型大小的同时保持模型的推理精度。针对该问题也出现了各种优化改进方法,例如量化误差补偿、自动量化及感知硬件量化等,这些都是热门的研究方向。 3.3.2感知量化训练 1. 算法介绍 感知量化训练方法最早由谷歌公司提出。它是一种伪量化训练,即在训练时模拟量化操作,过程中仍然采用浮点数计算,并通过反向传播学习更新网络参数,使得网络参数更好地适应量化带来的损失。由于在训练中已经引入了量化误差作为训练损失的一部分,优化器会在训练过程中尽量减少量化误差,从而得到更高的模型精度。 图3.36卷积层的伪量化训练图 一般地,模拟量化操作是通过在需要量化的参数后面插入伪量化节点来实现的。如图3.36所示,在卷积层(conv)的权重(weight)和激活层(ReLU6)后面分别插入权重的伪量化节点(wt quant)和激活的伪量化节点(act quant)。偏置(bias)后无须插入伪量化节点。当有多个卷积层时,上一层的输出(output)量化值即为下一层的输入(input)。 伪量化节点是感知量化训练中的重要因素,有以下作用: (1) 找到原始浮点数据的分布范围,即计算最小值rmin和最大值rmax。 (2) 通过“量化反量化”操作模拟量化过程中的精度损失。该损失会作用到网络模型中,传递给损失函数,从而影响到网络参数的更新。 根据前面的公式可推导出伪量化节点的正向传播公式 xout=s×cliproundxs+z,qmin,qmax-z (3.22) 式中,x是原始浮点数,xout是经过伪量化计算后的浮点数。由于round函数离散不可导,式(3.22)的反向传播梯度为0,导致网络权重无法在训练中正常更新。为了解决该问题,常用的一个方法是采用直通估计器(straight through estimator)方法来近似计算: 在伪量化节点的反向传播计算时,直接把前一层的梯度回传到伪量化节点之前的网络权重上,即忽略伪量化节点的影响。伪量化节点的正向传播和反向传播如图3.37所示。 图3.37伪量化节点的正向传播和反向传播 伪量化节点的反向梯度计算公式如下: Lx=Lxout×xoutx=Lxout,rmin≤x≤rmax 0,xrmax (3.23) 直通估计器方法解决了网络权重无法学习更新的问题,但同时也带来了梯度不准确的问题,而且随着量化比特数的降低问题愈加明显,这也导致了低比特量化网络很难进行高效训练。目前,业界已有针对量化梯度问题优化的一些研究。例如,在DSQ方法中引入可求导的tanh函数逐步逼近阶梯函数,目的就是为了缓解梯度不准确问题。 网络的感知量化训练完成后会保存量化模型,然后通过端侧推理框架在端侧设备上部署和推理。端侧推理框架会根据当前端侧设备的硬件配置进行硬件加速处理。其中,将归一化层(简称BN层)、ReLU层与卷积层融合是常用的一个推理加速技巧。因此,感知量化训练在对卷积层的权重做伪量化时,通常也会模拟推理时BN层和卷积层的融合,对融合后的权重参数做伪量化训练。 根据网络是否需要提前进行训练,感知量化训练可分为从头开始训练和在预训练好的模型上微调两种方式。一般来说,使用感知量化训练进行微调的效果会更好。算法3.1是在预训练好的模型上进行感知量化训练微调的流程。 算法3.1感知量化训练算法 输入: 预训练网络 输出: 量化网络 (1) 初始化: 设置权重和激活值的范围rmin和rmax的初始值。 (2) 构建模拟量化网络: 在需要量化的权重和激活值后面插入伪量化节点。 (3) 量化训练,重复执行以下步骤直至网络收敛。 ① 计算量化网络层的权重和激活值的范围rmin和rmax。 ② 前向计算,伪量化节点计算方式参见式(3.22)。 ③ 反向传播更新网络权重参数,伪量化节点的梯度计算参见式(3.23)。 (4) 导出量化网络。 ① 获取rmin和rmax,并计算量化参数s和z。 ② 根据式(3.14)计算权重的量化整数值,并替换对应网络层的参数和数据类型。 ③ 删除伪量化节点,在量化网络层前后分别插入量化和反量化算子。 感知量化训练也同样面临前面提到的量化精度挑战,如何更有效地训练量化网络仍是一个持续讨论的问题。IBM公司提出的PACT和LSQ方法将量化参数作为可学习参数,在量化到4bit时仍可以维持和原始精度模型类似的准确率。HAQ方法将强化学习和混合精度量化结合,自动找到适合网络层的量化比特数。TernaryBERT则是将知识蒸馏和2bit量化结合,在模型压缩为1/14.9的情况下可达到与全精度模型相当的精度。 2. 用MindSpore实现感知量化训练 感知量化训练与普通网络训练流程类似,主要有以下几个步骤: (1) 加载数据集,处理数据。 (2) 定义量化网络。 (3) 定义优化器和损失函数。 (4) 训练网络,保存模型文件。 (5) 加载保存的模型,进行推理。 (6) 导出量化模型。 其中,第(2)、(6)步是感知量化训练区别普通网络训练的操作。下面以LeNet网络为例,详细展开叙述第(2)、(6)中量化相关的操作。 使用MindSpore 感知量化训练的API前需要先导入MindSpore API和辅助模块,如代码3.13所示。 代码3.13导入MindSpore API和辅助模块 import mindspore.nn as nn from mindspore.compression.quant import QuantizationAwareTraining from mindspore import load_checkpoint, export, Tensor from mindspore.compression.quant import load_nonquant_param_into_quant_net from mindspore.compression.common import QuantDtype 首先定义量化网络。在生成量化网络之前,先将卷积层、全连接层和BN层、ReLU层写成融合网络层的形式,可以通过融合算子中的参数控制是否包含BN层和ReLU层,如代码3.14所示。 代码3.14替换融合网络层 class LeNet5(nn.Cell): def __init__(self, num_class=10): super(LeNet5, self).__init__() self.num_class = num_class self.conv1 = nn.Conv2dBnAct(1, 6, kernel_size=5, activation='relu') self.conv2 = nn.Conv2dBnAct(6, 16, kernel_size=5, activation='relu') self.fc1 = nn.DenseBnAct(16 * 5 * 5, 120, activation='relu') self.fc2 = nn.DenseBnAct(120, 84, activation='relu') self.fc3 = nn.DenseBnAct(84, self.num_class) self.max_pool2d = nn.MaxPool2d(kernel_size=2, stride=2) def construct(self, x): x = self.max_pool2d(self.conv1(x)) x = self.max_pool2d(self.conv2(x)) x = self.flattern(x) x = self.fc1(x) x = self.fc2(x) x = self.fc3(x) 代码3.14中融合算子Conv2dBnAct()中包含了原网络中的Conv2d和ReLU网络层的计算,融合算子DenseBnAct()中包含了原网络中的Dense和ReLU网络层的计算。 在预训练好的模型上进行感知量化训练微调,是基于预训练的权重进行更新学习,因此需要先加载预训练好的模型参数。 代码3.15加载预训练模型的参数 network = LeNet5(num_classes) param_dict = load_checkpoint(ckpt_path) load_nonquant_param_into_quant_net(network, param_dict) 接下来,调用转换接口生成量化网络。 代码3.16调用转换接口生成量化网络 quantizer = QuantizationAwareTraining(bn_fold=False, quant_dtype=[QuantDtype.INT8, QuantDtype.INT8], per_channel=[True, False], symmetric=[True, False]) net = quantizer.quantize(network) QuantizationAwareTraining. quantize()函数是量化网络的转换接口,同时也是量化策略配置的接口。bn_fold参数表示有BN层时是否与卷积层融合。quant_dtype参数表示权重和激活值的量化比特数,默认是8bit。per_channel参数表示权重和激活值是否逐通道量化,设置为False时表示逐层量化。symmetric参数表示权重和激活值是否对称量化,设置为False时表示非对称量化。 感知量化训练完成后导出在端侧硬件平台上可部署的量化模型。导出的量化模型中会删除伪量化节点和保存量化参数,然后通过端侧推理框架在端侧设备上部署和推理。 代码3.17导出量化后模型 inputs = Tensor(numpy.ones([1, 3, 32, 32]), mindspore.float32) export(network, inputs, file_name="mobilenet_quant", file_format='MINDIR', quant_mode='QUANT') 使用MindSpore对LeNet网络在MNIST数据集上进行感知量化训练微调,然后将导出的量化模型在MindSpore Lite推理框架上进行部署推理。 LeNet网络感知量化训练后测试结果如表3.2所示。 表3.2LeNet网络感知量化训练后测试结果 网络类型 模型大小/KB 推理准确率/% 原网络 245 98.45 8bit量化网络 66 98.65 从表3.2中数据可以看到,8bit量化后的LeNet模型大小约压缩为原来的27%,却仍然可以保持原始推理准确率水平。 MindSpore的感知量化训练功能除支持LeNet网络,还支持MobileNetV2、YOLOv3、BERT等网络的量化。在8bit量化的场景下,模型压缩2~4倍,量化后推理准确率无明显损失。MindSpore将逐步增强感知量化训练功能,例如支持不同任务的模型量化,并逐步支持低比特量化能力,以及与其他模型压缩方法结合,以满足开发者的不同使用需求。 3.3.3训练后量化 随着AI的发展,算力是不可绕开的话题。为了完成更具挑战性的任务,模型变得越来越复杂。而要在智能手机、手环、手表等功耗、内存、CPU受限的端侧、边缘侧设备上部署AI应用,更是挑战重重。MindSpore Lite致力于帮助用户打造极速、极智、极简的全场景智能应用,提供了深度优化的量化功能,能极大地压缩模型大小,同时提升推理速度,从而减小内存占用,降低功耗,使得训练好的模型能够轻松部署到端侧设备上。 神经网络量化有训练感知量化和训练后量化两种常用方法,训练感知量化在前面章节已介绍,本节主要介绍训练后量化。 1. 训练后量化方法简介 在训练后量化的具体实现上,有多种不同的方案。例如,对于卷积算子,可以整体量化权重参数,即逐层量化(perlayer),也可以逐通道量化(perchannel),即针对不同的通道,分别计算量化参数,提高量化的精度。一般而言,无论perchannel还是perlayer量化方案,对于权重的量化使用对称量化,对于激活值的量化使用非对称量化。其原因是,对于权重而言,数据的分布通常为对称的,如[-0.6,0.6]或[-8,8],因此采用对称量化可以进一步有效降低计算量,使得数据分布更加符合其真实的分布。在网络中通常使用ReLU和ReLU6等作为网络模型的激活函数,其数据分布集中在[0,6]或者[0,∞],如果采用对称方案会加剧数据的离散程度,另外会造成数据的浪费,如[-128,0]之间并不存储任何数据,只有[0,127]有数据分布。下面以8bit量化为例,介绍量化算法的主要细节。 1) 对称量化 对称的量化算法原始浮点精度数据与量化后INT8数据的转换如下: float=scale×int (3.24) 式中,scale默认是float32浮点数,为了能够表示正数和负数,int采用signed int8的数值类型。将float32原始高精度数据转换到signed int8数据的操作如下(其中round为取整函数,量化算法需要确定的数值即为常数scale): int=round1scale (3.25) 由于神经网络的结构以层(Layer)来进行划分,因此对权值和数据的量化可以层Layer为单位来进行,对每层的参数和数据分别进行不同scale的量化。 对权值和数据的量化可以归结为寻找scale的过程,由于int8为有符号数,要保证正负数值表示范围的对称性,因此对所有数据首先进行取绝对值的操作,使待量化数据的范围变换为[0,max],再来确定scale,其scale正确的计算公式为 scale=2max(|xmin|,xmax)Qmin-Qmax (3.26) offset的计算公式为 offset=0 (3.27) 确定了scale之后,itn8数据对应的表示范围为[-128×scale,127×scale],量化操作即为对待量化数据以[-128×scale,127×scale]进行饱和式截断,即超过范围的数据饱和到边界值,然后进行量化计算即可。 2) 非对称量化 非对称的量化算法与对称的量化算法主要区别在于数据转换的方式不同,同样需要确定scale与offset这两个常数,则有 float=scale×(uint+offset) (3.28) 确定后通过原始float32高精度数据计算得到uint8数据的转换,即 uint8=roundfloatscale-offset (3.29) 式中,scale是float32浮点数,uint是unsigned INT8定点数,offset是int32定点数,其表示的数据范围为[scale×offset,scale×(255+offset)]。若待量化数据的取值范围为[xmin,xmax],则scale的计算公式如下: scale=xmax-xminQmax-Qmin (3.30) offset的计算方式如下: offset=Qmin-roundxminscale (3.31) 对于权值和数据的量化,都采用上述公式的方案进行量化,xmin和xmax为待量化参数的最小值和最大值,Qmax和Qmin为 Qmax=2n,Qmin=0 (3.32) 3) 仅权重量化 仅量化模型的权重,在模型推理时,先将量化后的权重反量化为原始的fp32数据,后续的推理流程与普通的fp32模型一样。权重量化能压缩模型大小,精度误差较小。 权值在进行推理加速时均已确定,因此不需要提供额外的校准数据集。如果使用对称量化方案,max使用权值绝对值的最大值,对于非对称量化算法,max和min使用权值的最大值和最小值。根据算法验证结果,对于卷积,每个卷积核按照perchannel的方式采用一组独立的量化系数(scale和offset),量化后推理精度较高。因此,Convolution权值的量化根据卷积核数量分组进行,计算得到的scale和offset的数量与卷积核数量相同。Dense的权值使用一组scale和offset。 4) 全量化 全量化不仅会量化模型的权重还会量化模型的激活值,在模型推理时,会执行对应的量化算子,能够加快模型推理速度。不同于权重数据,在离线状态下,模型本身是不包含激活数据的,为了能够量化激活值,需要用户提供一定数量的校准数据集,用于统计每一层的激活值分布。校准数据集可以选自训练数据集(不需要包含标签)或者模拟真实场景中的输入数据,数量在100组左右即可。在做训练后量化时,会以校准数据集作为输入,执行推理流程,然后统计出每一层激活值的数据分布。利用数据统计到的每层激活值的数据分布,一般采用perlayer非对称的计算方式,便可以得到对应的量化参数。对于权重参数,直接根据权重的实际值,采用perchannel量化得到量化参数。 由于原始模型的输入、输出是fp32类型的,而量化后的模型中量化算子的输入、输出是int8类型的,为了数据类型连续,全量化模型会插入相应的量化算子(将fp32数据转成int8数据)与反量化算子(int8数据转成fp32数据)。 2. MindSpore Lite 训练后量化 通用的量化算法思路是将模型中的权重、激活值等float32类型的数据映射到int型的整数,如8bit量化,就是将神经网络中的float32数据映射到-128~127的整数区间,理论上能将模型大小压缩75%; 而在模型运行时,配合采用int8算子进行运算,相比较于float32类型的算子,计算速度能得到较大提升,功耗也能得到极大降低。但是,量化后的数据无法通过反量化无损地恢复出来,这就给模型带来了精度损失。为了提升量化模型精度,MindSpore Lite实现了多种激活值量化方案及BiasCorrecton、Mean/Variance Correction精度校正方法,使得模型量化精度远远领先于业内其他框架。 在对神经网络进行量化时,以卷积算子为例: 由于权重的值是固定、事先可获取的,MindSpore Lite直接采用线性非对称量化。q=round(r/S+Z),q表示量化后的数据,r是float32数据; S和Z是量化参数scale、zeropoint,S=(r_max-r_min)/(q_max-q_min),Z=round(q_max-r_max/S)。由于权重的各个通道(channel)是相互独立的,对不同的通道分别进行量化能够显著提高精度,MindSpore Lite便采用这种PerChannel的方法。不同于权重,激活值在量化时无法获取,通常在模型转换阶段提供校准数据集,通过统计校准数据集的数据,来计算量化参数。但是校准数据集无法完整地表征实际推理时的输入,所以激活值量化会带来较大误差。为了保障训练后量化的精度,MindSpore Lite提供了以下三种激活值量化策略。 1) MAX_MIN(最大最小策略) 直接将激活值数据从[-|max|,|max|]映射到int8区间,从而计算出量化参数。对于校准数据集能够比较好地表征实际输入数据集的情况,本方法就已经足够了。MAX_MIN量化如图3.38所示。 2) RemoveOutlier(移除离群点策略) 考虑到校准数据集的数据可能分布不均匀,例如绝大多数数据分布在[-1,1]这个区间,只有极少的数据在[-1,1]区间外,那么便可以直接将[-1,1]映射到int8区间,可以获得更高精度。这些极少的分布在[-1,1]区间之外的数据称为离群点(outlier),将这些离群点剔除,再进行激活值量化,能够提高模型精度。RemoveOutlier方法将数据按序排列后,只保留在中间部位一定比例的数据(默认90%),而将首尾部分的较小值、较大值剔除,如图3.39所示。 图3.38MAX_MIN量化 图3.39离群点移除 3) KL策略 与RemoveOutlier方法类似,KL策略通过选择合适的区间[-|T|,|T|]映射到int8,来提高模型精度。不同的是,为了找到最合适的T,KL方法通过搜索的方式,先将原始数据的分布统计成直方图; 然后将数据从T处截断,经过T截段后的数据再重新统计直方图,与原始数据的分布直方图进行对比,计算KL散度,选择KL散度最小的T,也就是数据分布相似度最高的截断方式,作为最终的阈值,从而获得区间[-|T|,|T|],使得量化误差极小化,如图3.40所示。 图3.40KL散度 该方案的思路为最小化量化后数据的统计分布与原始高精度数据的统计分布的差异性,操作流程如下: (1) 使用直方图统计的方式得到原始float32数据的直方图统计分布Pf。 (2) 在给定的min和max搜索空间中选取若干个minq和maxq分别对待量化数据进行量化,分别得到量化后的数据Qq。 (3) 使用同样的直方图统计的方式得到n个Qq的直方图统计分布Q。 (4) 分别计算Q中每个Qq与Pf的统计分布差异性,找到差异性最低的一个对应的min和max作为确定的量化数值。 在上述操作中,涉及的超参数包括进行直方图统计时选取的直方图bin个数、min和max的搜索空间、统计分布差异性的指标。 直方图bin个数直接反映了直方图统计特征的分布数个数,由于数据经过量化后会集中到256个离散的点上,因此bin的个数不宜过大,否则绝大多数的bin都没有数值。 max的搜索空间可以通过search_start_scale、search_end_scale与search_step来确定。 (1) search_start_scale为搜索空间起始点与search_value的比值。 (2) search_end_scale为搜索空间结束点与search_value的比值。 (3) search_step为搜索空间中每次搜索的值与seach_value的比值步进值。 以max_candidate=100,search_start_scale=0.8,search_end_scale=1.2,search_step=0.01为例,对称量化算法下,其定义的max搜索空间为从100×0.8=80到100×1.2=120的范围,每次步进100×0.01=1,一共41个d_max搜索值; 非对称量化算法下,搜索0.8×(max-min)~1.2×(max-min),确定最好的一个系数。 另外一个例子,search_start_scale =0.3,search_step=0.01,search_end_scale==1.7,bin=150。需要在0.3×max~1.7×max这个范围中找一个最好的max,搜索步长是0.01max,因此需要搜索(1.7-0.3)/0.01+1=141个max数值。直方统计的bin个数也可以设置,假设当前是150。对于算法方案一,将0~2×max的数据分为150段,统计数据落在每段中的频率,使用数据的绝对值统计; 对于算法方案二,在0.3×(max-min)~1.7×(max-min)这个范围中找一个最好的ratio,将0~2×(max-min)的数据分为150段,统计数据落在每段中的频率,使用datamin的值来统计频率。141个数值都要统计,因此一个量化算子需要存储141×150个数值。多个batch情况下对frequancy进行累加。 统计分布差异性的指标为计算两个长度为n直方图Pf、Qq分布之间的信息差异度,选取的指标包括以下三种。 (1) KullbackLeibler Divergence(KL散度),计算公式如下: DKL(Pf‖Qq)=∑Ni=1P(i)×log2Pf(i)Qq(i) (3.33) (2) Symmetric KullbackLeibler Divergence(对称KL散度),对于Pf、Qq两个分布,Pf相对Qq与Qq相对Pf的KL散度是不同的,对称KL散度的计算方式为Pf相对Qq与Qq相对Pf的KL散度的均值,计算公式如下: SYMKL=12×DKLPf‖Qq+DKLQq‖Pf =12×∑ni=1Pf(i)×log2Pf(i)Qq(i)+∑ni=1Qq(i)×log2Qq(i)Pf(i) (3.34) (3) JensenShannon Divergence(JS散度),首先生成一个新的分布M为Pf与Qq的均值,JS散度为Pf相对M的KL散度与Qq相对M的KL散度的均值,计算公式如下: M=12×(Pf+Qq) (3.35) DJS=12×DKLPf‖M+DKLQq‖M =12×∑ni=1P(i)×log2Pf(i)M(i)+∑ni=1Q(i)×log2Qq(i)M(i) (3.36) 为了消除量化的固有误差,MindSpore Lite采用了事后修正的方法,对量化误差进行校正。以卷积为例,a=∑Niwixi+b,w表示卷积的权重输入、x表示激活值输入,b是偏差(bias)输入。经过量化卷积算子计算出来的结果记为q_a,对量化算子的输出结果进行反量化得到float32类型数据a^,由于量化天然存在误差,a^与a的值相比存在误差Δ,a^=a+Δ。MindSpore Lite巧妙地将误差Δ补偿到bias输入上去,从而降低量化误差。具体做法是: 在模型正常完成量化后,逐层计算量化模型的输出,并与FP32模型的输出进行比对,计算得到误差Δ; 将Δ按照channel取平均,加到卷积的bias参数上,从而抵消量化模型输出误差,即b=b-Δ。训练后量化中误差校正如图3.41所示。 图3.41训练后量化中误差校正 同时注意到,量化引发的误差除体现在数值不匹配外,还体现在统计特性上。量化前权重数据的均值、方差分别记为E(w_c )、‖w_c-E(w_c)‖; 而量化后的数据经反量化,其均值、方差会发生变化,分别记为E((w_c )^)、‖(w_c )^-E((w_c )^)‖。由此,MindSpore Lite对量化数据进行均值/方差校正。首先计算根据式(3.37)和式(3.38)计算出均值校正因子,然后根据式(3.39)计算方差校正因子,最后便可以对权重的逐通道数据进行校正。 w^c←ζc(w^c+uc)(3.37) uc=Ewc-E(w^c) (3.38) ζc=‖wc-E(wc)‖‖w^c-E(w^c)‖ (3.39) 除在精度上的极致追求外,MindSpore Lite也支持1~16bit的任意比特量化,为用户提供更多的压缩率选择,同时正在积极完成混合比特量化等新特性,致力于为用户提供极速、极智、极简的AI框架。 3. MindSpore Lite训练后量化实践 MindSpore Lite训练后量化功能集成在converter_lite工具中,具体的使用方式请参考https://www.mindspore.cn/tutorial/lite/zhCN/master/use/post_training_quantization.html。大量的实验及实际应用的交付表明,MindSpore Lite的量化功能够非常好地实现模型压缩,降低功耗,同时具备很好的精度。如表3.3所示,以mobilenet_v2模型为例,MindSpore Lite训练后量化A8W8(激活值8bit量化、权重8bit量化)精度与FP32模型相比只损失0.4%。 表3.3MindSpore Lite训练后量化精度结果 mobilenet_v2 精度(在ImageNet) FP32 71.56% A8W8 71.16% A8W8(不进行bias_correction) 70.74% A8W7 71.06% A7W7 70.78% 表3.4展示了是否使用Correction对量化精度误差的影响,可以看见,基于误差的校正手段能够显著提高量化模型的精度。 表3.4BiasCorrection精度对比结果 模型 不使用Correction 使用Correction ml_face_tracking 1.28214% 1.00386% ml_face_age 8.40158% 3.21885% ml_face_beard 3.72371% 1.50271% ml_face_contour 0.196258% 0.129654% ml_face_emotion 17.5848% 3.19592% ml_face_glasses 5.8514% 4.73312% 3.3.4小结 MindSpore同时具备感知量化训练与训练后量化的功能。对于需要重新训练模型的场景,可以使用感知量化训练获取压缩模型,能够具备非常高的精度; 如果无法重新训练模型,可以使用训练后量化,快速地压缩并部署模型。 3.4类型推导 在深度学习领域中,深度学习框架一般使用Python语言作为描述网络模型结构的前端语言。Python语言是一门动态弱类型的语言,入门简单,开发代码简洁高效,但由于其解释执行,运行速度慢。Python前端语言给用户带来了动态灵活的语义和高效的开发效率,但是想要生成运行高效的后端代码,需要给后端框架提供优化友好的静态强类型中间表示(Intermediate Representation,IR)。因此,需要一种高效可靠的静态分析方法为桥梁,将Python前端表示转换成等价的静态强类型IR,同时给用户带来高效的开发效率和运行效率。 HindleyMilner(HM)类型系统是一种经典的具有参数多态性的简单类型lambda演算的类型系统,也被称为DamasMilner或DamasHindleyMilner。它最初由J. Roger Hindley 描述,后来由Robin Milner重新发现。路易斯·达马斯(Luis Damas)在其博士学位论文中对此方法进行了详尽的形式分析和证明。因为算法中使用的过程的成本接近O(1),所以算法的总成本在要推断类型的表达式的大小上接近线性。HM最显著的特性是它的完整性和推断给定程序的最一般类型的能力,无须程序员提供类型注释或其他提示,而且性能高。 MindSpore静态分析模块借鉴了Myia的Infer模块。Myia的Infer模块采用抽象释义(abstract interpreter)的方式对IR中的节点进行抽象值(包括具体值、类型、维度)的推导。为了支持控制流节点的推导,Myia引入协程来同时执行多个分支的推导,以解决“无法知道哪个分支是递归的退出条件”而可能导致推导进入死循环的问题。但是引入协程后函数执行顺序不确定,不利于程序的跟踪调试,同时协程机制会传到整个Infer模块,导致处理逻辑复杂。 3.4.1静态分析技术背景 1. 静态分析常用术语 (1) 静态分析方法: 静态分析是指在不实际运行程序的情况下,对代码进行分析的技术。 (2) 控制流图: 控制流图是反映程序执行过程的图结构抽象表示,图中每个节点表示一个基本代码块,其中没有分支跳转等控制语句,每条有向边表示一个分支跳转,图中的每条路径反映了程序可能的实际执行过程。 (3) 控制流分析: 控制流分析是分析程序控制流图、获得程序静态属性的静态分析方法。常见的控制流分析方法包括抽象释义、约束补偿和类型系统等。 (4) 抽象释义: 抽象释义简单来说是通过抽象解释器将语言的实际值、实际语义近似为抽象值、抽象语义进行不确定性的解释执行,从而来获得数据流分析想得到的属性(即抽象值中所承载的信息)。抽象值一般包括变量的类型和维度等。 (5) IR: IR是深度学习框架分析前端代码转换成的等效的通用内部表示形式,以方便进一步的分析和优化流程。 (6) 函数式IR: IR中的IF、WHILE等分支控制语句都通过函数定义和函数调用来表示,每个分支跳转都对应了一个高阶函数调用: 输入分支判断条件和两个分支函数,输出对应分支的执行函数。调用高阶函数的输出函数,可得到分支跳转的结果。 2. MindIR MindIR是基于图表示的函数式IR。MindIR文法继承于ANF文法,ANF 文法定义如下: ::= NUMBER | STRING | VAR | BOOLEAN | PRIMOP |(lambda (VAR …) ) ::= ( …) |(if ) ::= (let ([VAR ]) ) | | ANF中表达式分为原子表达式和复合表达式,原子表达式表示一个常数值或一个变量或一个匿名函数; 复合表达式由多个原子表达式复合组成,表示一个匿名函数或原语函数调用,组合的第一个输入是调用的函数,其余输入是调用的参数。 MindIR文法定义如下: ::= | ::= Parameter ::= Scalar | Named | Tensor | Type | Shape | Primitive | MetaFuncGraph | FuncGraph ::= ( …) ::= | MindIR中的ANode对应于ANF的原子表达式,ANode有两个子类分别为ValueNode和ParameterNode。ValueNode表示常数节点,可承载一个常数值(标量、符号、张量、类型、维度等),也可以是一个原语函数(Primitive)或一个元函数(MetaFuncGraph)或一个普通函数(FuncGraph),因为在函数式编程中函数定义本身也是一个值。ParameterNode是参数节点,表示函数的形参。MindIR中CNode对应于ANF的复合表达式,表示一次函数调用。 3. 控制流用例场景 控制流在MindIR中是以高阶函数选择调用的形式表达。高阶函数选择的形式,可以支持任意的条件跳转、循环和递归等控制流场景。一个典型的控制流场景是Fibonacci,如代码3.18所示。 代码3.18Fibonacci控制流示例 @ms_function def fibonacci(n): if(n < 1): return 0 elif(n == 1): return 1 else: return fibonacci(n-1) + fibonacci(n-2) 对应的MindIR如图3.42所示。 图3.42Fibonacci函数MindIR 其中fibonacci是顶层函数图,在顶层中有两个函数图被switch选择调用。√fibonacci是第一个if的True分支,×fibonacci是第一个if的False分支。在×fibonacci中被调用的√×fibonacci是elif的True分支,××fibonacci是elif的False分支。这里需要理解的关键是,在MindIR中,条件跳转和递归是以高阶函数的形式表达的。例如,√fibonacci和×fibonacci作为switch算子的参数传入,switch根据条件参数选择哪一个函数作为返回值。因此,switch是把输入的函数当成普通的值做了一个二元选择操作,并没有调用,而真正的函数调用是在紧随switch后的CNode上完成的。 上述的MindIR是经过了静态分析得到的静态强类型的IR,其中Python定义里的加法操作被特化为scalar_add常数加法操作。静态分析的目标就是分析初始MindIR中每个节点的类型和维度等,对函数和原语进行特化。 3.4.2静态分析设计 1. 总体设计 MindSpore静态分析模块位于MindSpore框架前端,是前端优化中的一个环节,通过前端提供的IR遍历和克隆等工具,将MindIR转换成静态强类型的IR。静态分析模块任务总流程如图3.43所示。 图3.43静态分析任务总流程 2. 静态分析架构 静态分析模块是执行抽象释义方法的抽象解释器,用于分析MindIR中节点的抽象值。MindSpore采用的方法是ADI方法的演进,符合抽象释义的特征: 对抽象值做不确定的抽象语义的解释执行,停机后每个节点的抽象值是所期望得到的程序静态信息。基本的抽象释义方法流程可以理解为从MindIR的顶层函数图入口开始解释执行,将函数图中所有节点进行拓扑排序,从前往后根据节点的语义递归推导各节点的抽象值,遇到函数子图时,递归进入函数子图进行解释执行,最后返回顶层函数return节点的抽象值。根据抽象释义方法流程,静态分析模块主要涉及抽象值类型系统的设计、抽象值缓存机制的设计、各语义推导方法的注册及算法停机保证的设计。 静态分析模块主要分为抽象域模块、缓存模块、控制流处理算法模块和语义推导模块,如图3.44所示。 图3.44静态分析模块 3.4.3静态分析模块设计 1. 抽象域设计 抽象释义方法执行过程中需要将语言中的实际值近似为抽象值。一般来说,实际值与抽象值之间需要满足伽罗瓦(Galois)连接,抽象值需要设计为完备格以满足有限单调性系统设计,从而保证不确定性的近似语义可以停机。静态分析系统的抽象域信息主要包含类型、维度和具体值。MindIR中每个节点类型都有对应的抽象值。例如,MindIR中Number的抽象值是AbstractScalar(type,value),Function的抽象值是AbstractFunction。静态分析模块中所有抽象值类都继承AbstractBase类,如图3.45所示。 图3.45抽象值类继承关系 抽象域中具体值和类型会有顶层的抽象,如kAnyValue和kAnyType,表示任意具体值和任意类型。静态分析过程中,为减少顶层函数图的编译,让其能够处理更多的抽象值类型,会将其参数对应的抽象值中的具体值设置为顶层抽象kAnyValue。抽象值有两个操作,泛化(Broaden)和合并(Join)。泛化是指将节点的抽象值中的具体值提升成顶层的抽象; 合并是指将两个抽象值合并,合并过程中,如果两个抽象值不一致,会对抽象值的具体值进行泛化,提升成顶层的抽象,如图3.46所示。 图3.46抽象值Broaden和Join例子 2. 支持上下文的缓存机制 抽象释义的缓存机制支持上下文环境(context),也就是抽象解释方法会维护函数被调用时的上下文环境。通过函数和函数被调用时的上下文环境捕获的自由变量构成的闭包(closure),可以唯一确定函数内节点的抽象值。解释执行过程中,推导过的节点上下文环境和抽象值都会记录在缓存中。在抽象解释执行过程中遇到具有相同上下文环境的函数图节点,会从缓存里读取该节点的抽象值,如图3.47所示。 图3.47缓存机制 基本的缓存机制如图3.47所示。AnalysisEngine类是执行抽象释义方法的引擎。AnalysisCache类包含一个映射unordered_map,维护计算过的节点抽象值。AnfNodeConfig类保存节点和节点的上下文环境信息。AnalysisContext维护节点的上下文环境,主要记录节点所属函数子图、函数参数的抽象值及上层函数调用的上下文环境。EvalResult类包含节点抽象值和属性。AnalysisEngine类方法Run函数的输出是AnalysisResult类。AnalysisResult类包含顶层函数子图的返回值节点的抽象值信息和上下文环境信息。AnalysisEngine类Run函数执行过程中会调用Evaluator函数推导类,Evaluator类也包含一个cache映射unordered_map,保存该函数或函数闭包在不同参数下返回值的抽象值。 3. 各语义的推导方法 表达式语义主要基于MindIR的文法定义,这些语义的推导方法主要靠Evaluator类实现。Evaluator类提供了Run方法,输入算子或函数参数的抽象值,输出算子或函数输出的抽象值。具体的语义推导方法参见图3.48 Evaluator类继承关系。 各Evaluator类的作用参见表3.5。 表3.5Evaluator类的作用 Evaluator类 作用 PrimEvaluator 负责普通算子原语的推导 BaseFuncGraphEvaluator 负责FuncGraph、MetaFuncGraph函数子图的推导 TrackedEvaluator 对Primitive和FuncGraph Evaluator封装,用于区分对同一个函数图或者Primitive的不同调用点的推导 PartialAppEvaluator 负责对偏函数调用的推导 VirtualEvaluator 负责有固定函数签名和抽象值的函数的推导 PartialEvaluator 负责偏函数算子的推导 JEvaluator 负责Grad算子的推导 MixedPrecisionCastEvaluator 负责混合精度转换算子的推导 UnpackGraphEvaluator 用于子图可变参数展开 DoSignatureEvaluator 用于算子类型转换 这里以Python算子的推导方法注册为例进行简单介绍。Python算子的抽象值推导方法由PythonPrimEvaluator类进行调用。注册Python算子的推导方法只需要继承PrimitiveWithInfer类,并重写PrimitiveWithInfer类的infer_dtype、infer_shape、infer_value或__infer__方法。图3.49是Log算子的推导方法的注册。 4. 控制流处理算法 1) 不确定性语义处理设计 不确定性语义是指解释执行过程中遇到if或while等分支控制流语句时,若分支条件不确定,也就是条件中变量的具体值未知时,不能确定下一步应该执行哪一个分支。这时算法会解释执行所有分支,并将所有分支返回的抽象值合并作为整个语句的抽象值。将上 图3.48Evaluator类继承关系 图3.49Log算子的推导方法的注册 述分支称为不确定性分支。 由于我们会对循环的所有分支进行推导,当循环过程中存在不确定性分支时,循环体函数可能会被重复推导,也就是经过循环体的解释执行,循环入口条件某一变量值发生改变,导致每次循环体捕获的context均不一致,从而造成死循环。 因此,静态分析方法检测到不确定性语义和循环调用结构时,会判断循环入口条件变量的抽象值是否发生变化。若发生变化,将入口条件变量抽象值的具体值泛化成顶层抽象值,避免无限的推导。一个简单的例子如图3.50所示。 图3.50不确定性语义处理机制示例 2) 循环嵌套处理设计 抽象解释执行过程中会遇到循环嵌套的结构,为了退出循环并保证每一条语句都正确推导出抽象返回值,需要遍历每一个循环出口。如图3.51所示的片段代码,内层的×fac分支调用的分支中又包含了一个嵌套的fac调用。 图3.51不确定性语义处理机制 结合之前提到的不确定性分支,循环出口是指已经推导的某个不确定分支匹配的尚未推导的不确定分支,比如可能是上述√×fac分支对应的××fac分支。为了处理循环嵌套,使用两个数据结构进行相关信息的记录。 (1) 使用multipossibility trace (list>)记录解释执行过程中已经推导过的不确定分支。 (2) 使用multipossibility map映射(map)维护不确定分支的互相映射关系,如(√×fac分支>××fac,××fac>√×fac分支)。 根据不确定性分支是否在trace中,将不确定性分支的执行方式分为direct run和trace run: direct run方式是直接执行分支,但是不将分支添加到trace中; trace run方式是执行时将不确定性分支添加进trace中。算法的基本思路是: 如果当前正在推导的分支及相对应的分支都在trace中,那么就从当前推导的分支开始,查找在trace中还没有执行的子分支,如果存在,则那个分支就可能是循环出口,使用direct run的方式进行推导,直到循环出口返回,如图3.52所示。 图3.52不确定性出口及循环嵌套处理机制 3.4.4小结 本节描述了通过抽象释义的控制流静态分析技术将Python前端网络转换成静态强类型MindIR。一方面,满足用户使用Python语言作为描述网络模型结构的前端语言,用户可以简易地使用Python原生的控制流语句,编写包含循环和递归调用等复杂的控制流场景,根据需求可以高效地进行复杂网络模型的设计和运行。另一方面,静态强类型MindIR 为后端优化提供基础。 当前的控制流处理算法还有不完善的地方,对于某些控制流还无法处理,需要进行完善。一种可能的思路是对控制流相关的图采用两遍推导的方式,通过图分析,找出可能的出口分支,第一遍对从函数图入口到函数图出口的路径的子函数图进行推导,第二遍对全函数图进行推导。 3.5图算融合 随着AI的发展,深度神经网络模型结构快速演化,如GPT3等大模型(GPT3 有2048个隐藏层单元)越来越多,对运行性能和效率的要求也越来越高。人工手工优化的方式,愈发无法满足日益增长的需求。为此,业界AI框架进行了大量的探索,并取得了优秀的成果。 XLA(Accelerated Linear Algebra)是Google公司在2017年推出的TensorFlow编译器。XLA使用JIT编译技术分析计算图,将其转换为高级优化器(High Level Optimizer,HLO)中间表示(Intermediate Representation),在更细粒度上挖掘更大的图优化空间,通过与目标无关的优化和分析过程、与目标无关的运算融合,以及用于为计算分配运行时内存的缓冲区分析来实现收益,并最终通过低级虚拟机(Low Level Virtual Machine,LLVM)自动生成对应的机器代码。在2020年7月29日,Tensorflow通过XLA 在BERT MLPerf 提交的任务中,用 8 个 Volta V100 GPU实现了约7倍的性能提升。 TVM是陈天奇基于Halide提出的深度学习自动代码生成工具,TVM提供了Relay高层图表示,并为此提供了包括操作融合在内的图优化能力,而在代码生成时,TVM将计算和调度进行抽象分离,同样的计算逻辑可以针对不同的硬件进行调度处理,在不同后端上自动生成对应的高效执行代码。 从XLA和TVM等业界趋势来看,基于自动算子编译的计算图层面(下文称图层)和算子层面联动融合优化是提高模型运行性能和效率的重要手段。图算融合正是从这点出发,在MindSpore框架中联合图层和算子层进行深度优化,助力模型在MindSpore中获得极致性能。 3.5.1技术原理 优化计算图的执行性能,就是让特定计算图经过系列转换处理后,在对应硬件平台上执行得更有效率,充分利用已有的算力资源来完成目标计算任务。按阶段分,主要分为图层变换优化和算子优化。图算融合,目的是联合图层与算子两个阶段,让图层阶段计算图结构的变换更有利于高效算子的生成,同时让算子生成的规律反向指导图层变换的进行,联动起来实现计算图极致的性能表现。图算融合主要从算子优化、数据访存优化、硬件使用率、计算逻辑优化等多个角度进行了思考。 1. 算子融合优化 1) 访存融合 构成计算图的主要单元是一个个代表具体计算逻辑的算子节点。算子节点间通过数据或控制关系相互联系而成图,而这些各自独立的算子节点大多数对应有一份独有的可执行逻辑,并在图执行时分别被调用以承担独立的计算任务。 算子间通过数据或控制关系联系起来,这种关系为每个算子产生了独立的边界,造成了计算的孤岛。边界的存在保证了算子的计算完整性和计算独立性,也增加了相对应的数据访存需求: 一方面,计算完整性与独立性让针对性的算子优化得以实现,但同时也舍弃了更大范围的算子优化可能; 另一方面,如果算子的输入和输出数据的数据量比较大,则访存在执行总耗时上占比高,降低了计算访存比,让数据访存成为性能上的瓶颈。图3.53是NVIDIA P100、V100、A100 三代硬件在计算和访存吞吐率上的改进情况(其中,TFLOPS为Floatingpoint operations per second的缩写; FP32为Float32的缩写; FP16为Float16的缩写; Bandwidth指带宽): 计算和访存吞吐率得到了大幅提升,与之相比,访存带宽则整体变化不明显。因此,通过增大计算的密度,提高计算有效率,同时减少访存的需求,便可以合理利用硬件的发展趋势,完成性能的进一步优化。通过访存关系对前后算子进行融合,可以有效地减少访存数,提高计算密度,是一种重要的优化手段。 图3.53NVIDIA 三代硬件计算和访存吞吐率上的改进情况 计算图中的算子节点,主要有三类: (1) 不可再分的基础计算(basic op)。如加减乘除。在构建深度学习网络时,通过基础算子原语来表征计算,可以让计算图的构建更为灵活多样,不受限于高层接口API,但会给计算图引入大量细小的基本计算,造成较差的计算访存比。 (2) 具有特定含义的复合算子(composite op),如卷积、矩阵乘法、Softmax。卷积和矩阵乘法操作这类复合算子计算密度极高,为了提高其计算速度甚至设计出了专用的硬件单元(如NVIDIA Tesla V100引入的TensorCore),这种算子通常可以保留独自存在,通过调用硬件方提供的接口完成硬件上的特殊优化,如调用cuBLAS/cuDNN等库。而对于Softmax、ReLU、BatchNorm这些复合算子,虽然也有相应的高效接口,但其本身的计算上存在优化空间,其边界是需要重新考虑的。 (3) 通过分析优化而固化下来的复合算子。这类算子常出现在计算图的图层变换过程中。在工程中,经人工经验总结出融合模式,并在此基础上进行算子的优化定制,获得匹配的收益。这类算子在一定程度上缓解了固定算子边界的局限性,并提供了一种人工定制的方式。但由于需要精准的匹配模式,其适用的场景较为限制,而且如果需要匹配更多的场景,则对应地需要罗列更多的模式,使总算子量膨胀式增长。同时,算子的优化定制对编写算子逻辑的具体实施者能力要求十分高,其开发成本也是巨大的。 针对以上三类算子,图算融合通过吸纳整合基础算子,择劣处理复合算子,实现更优的访存融合优化,总体如图3.54所示。 图3.54算子边界的破而后立 在计算图将存在优化空间的复合算子摘取出来,并将其通过扩展器(Expander)一一扩展为子图,如图3.54(a)中围住一个或多个扩展出来的算子(expanded op)虚线所示。目标复合算子可以通过统计的方式,在已存的深度学习网络库中进行测试,适者留存不变,劣者则被选中处理; 也可以通过试错方式,在目标优化网络中先进行尝试,以结果作为选择标准。复合算子对应的子图由更细粒度的算子以符合其原始语义的方式组合而成,新的子图将在计算图中替代原始的复合算子的位置。子图为计算图增加了一种新的边界,这种边界本身是弱边界,其内部计算逻辑是可见的,为进一步优化提供了可能性。 为了提高计算密度,减少访存次数,一种简单的思路是: 让算子边界所能表达的计算逻辑更多,算子边界的数量更少。图算融合从这种思路出发,先尽可能多地去融合算子,同时以合理有效的方式对融合后的巨大边界进行拆解,重新划分出新的边界。从访存角度出发,各自相邻的算子或子图间通过融合的方式消除其间的边界。融合时需要保证融合后的大边界不能给计算图造成环结构(如图3.55所示,A与C融合将形成局部环),形成环结构的各个部分间由于互相依赖,将无法得到确信的执行时序,让整个计算图不可执行。对整个计算图进行访存融合后,边界将被推进到不被进行扩展的复合算子边缘,形成亟待优化的庞大计算语料块。但融合出来的计算语料本身是粗放的,直接编译执行并不能保证计算的高效,甚至其中造成的执行性能劣化可能会超过因访存减少带来的收益。因此,合理地设计新边界,可以通过必要的访存损耗,赢取更多的优化机会。 图3.55错误融合出现的环 如何平衡计算与访存之间的损益关系,划分合理边界,并让生成的算子质量更高?自动内核生成器(Auto Kernel Generator,AKG)针对该问题交出了一份令人满意的答卷。AKG基于多面体(polyhedral)技术实现了优质的自动调度,其中包括自动向量化、自动切分、线程/块(thread/block)映射、依赖分析和数据搬移等,并针对后端进行包括TensorCore使能、双缓冲区、内存展开和同步指令插入在内的各项优化。给定具体的算子描述,AKG将通过系列分析生成优质的可执行算子。AKG为图算融合提供了自动算子生成的能力,也为算子性能的优质性提供了保障。 但算子优化是有极限的,更优的算子需要图层变换的协作配合。AKG一方面保证了算子的性能,另一方面AKG生成算子的经验和对后端硬件平台优化的思考,也在图层重新划分算子边界中起到导向性作用。若将粗放式融合后的计算语料中的细粒度计算逻辑分成多种类型,如逐元素(elemwise)类型、广播(broadcast)类型、归约(reduce)类型等,针对这些类型设计成块规则(block rules)或在此基础上形成更抽象的代价模型(cost model),并将其应用于划分边界的评估中来,便可以得到最终需要的各个融合算子,并在AKG中完成高效算子的生成。 图算融合针对算子边界形成的访存和计算优化问题,对计算图的算子进行先聚合、再拆分,以这种破而后立的方式,让计算优化可以跨越原始算子边界进行,降低了单一计算逻辑带来的优化局限性的影响,并且不一味地扩大聚合范围,依据评估结果生成收益更高的必要边界,让访存成为提升性能的手段。同时AKG自动生成高效可执行算子的能力,让图算融合突破了传统一对一的手工算子融合的限制,让融合的形式更灵活有效。图算融合的访存融合使计算图中算子的计算密度及计算有效性得到有效提升,同时减少了算子边界间低效的访存操作,进而改善全图执行性能。 2) 并行融合 为了让计算图能够正确地执行,通常需要为各个计算任务分配执行顺序,计算任务也因此被隐式地添加上了执行上的依赖关系。最差情况下,在某一时刻,设备将被一个计算任务所占有,该计算任务将决定设备在当前时刻的利用率。如果大量细碎的计算任务被加载到设备中执行,则将可能让设备在大多数情况下处于空闲状态,硬件没有得到完全使能,造成计算资源的浪费。为此,图算融合提出了并行融合: 分析计算图中算子间的依赖关系,将存在并行性质且对硬件资源使用率低的算子融合在一起。在生成该并行算子时,对各部分独立分析调度,并安排相互独立的计算资源。如图3.56所示,其中Op A和Op B均为细小算子,正常执行时不能用满硬件资源。Op A和Op B单独存在时,将会在不同的阶段被加载到设备上执行,分别占用一部分执行时间; 将Op A和Op B进行并行融合后,融合后的算子被加载,设备将得以在足够的计算资源下同时执行属于Op A和Op B的计算逻辑,这样,既降低了运行时设备的空闲率,也让整体的运行耗时变短(执行时间由Op A加上Op B的耗时,变为Op A和Op B两者耗时的较大值)。 图3.56并行融合算子运行 并行融合在保留生成算子优质性的同时,将计算资源使用率小的算子捆绑成一个大算子,提高设备的使用率,减少因空置造成的计算资源浪费,提升整体执行性能。 2. 跨边界算子计算优化 粗放式聚合起来的计算逻辑状态并不是最优的,在进一步拆解前需要进行计算逻辑上的优化。聚合前存在一些优化壁垒,例如: 在未聚合时,由于算子边界的存在,算子只能完整引入,因此在构建深度学习网络时,为了其中部分输出的结果而引用了复合算子,会将其中不需要的冗余操作也一并带入网络中来; 不同复合算子的操作逻辑间含有相似的操作,在未进行聚合前,这些操作逻辑都是不可分割合并的。这些问题,随着计算逻辑的裸露和聚合,相互间的关联变得可见,优化便变得可能。在聚合起来的计算块中,通过公共子表达式消除(Common Subexpression Elimination,CSE)、死代码消除(Dead Code Elimination,DCE)等手段可以提高计算逻辑有效率,减少无效操作。 1) 公共子表达式消除 公共子表达式消除是在计算逻辑中搜索出存在被多次运算的部分操作,分析并考虑通过一个临时变量来保存中间结果的过程。如图3.57(a)所示,b和c的乘法操作存在两次,其结果分别和a相加、和d相乘。通过CSE处理后,b和c的乘法操作合并成一个,临时存储中间结果。CSE操作后,整体上少了一次乘法操作。 2) 死代码消除 死代码消除是将计算逻辑中对最终输出结果不造成影响的部分操作移除。如图3.57(b)所示,原计算逻辑中,对c取反并将结果与d相乘的操作并没有被外部使用,其本身是冗余操作,不对整体结果造成影响。通过DCE处理后,计算逻辑被精简,一方面减少了冗余操作,提高了计算有效率; 另一方面,如果消除的部分包含输入、输出,还可以减少算子的访存需求。 图3.57CSE和DCE 由于需要重建新算子边界,不经过处理还会将优化机会掩盖在新的算子中,形成新的局部性。例如: 聚合在一起的算子中,可能存在如同图3.58(a)的结构,该结构中包含两个轴求和(ReduceSum,一般需要指定求归约的轴axis参数)操作,如果不加处理,则将被划分为图3.58(b)三个算子,新建两个边界。其中取反(Negative)操作不影响轴求和过程,但由于它的存在隔开了两个轴求和。如果采用代数化简的方法,将取反操作下移到轴求和后,并合并两轴求和,结构将变换为图3.58(c)所示结构,变换后重新划分为两个算子,减少了一个算子,而且算子间的访存量也更小。从这个例子可以看出,在相同的计算量下,通过一些恒等的变换(如代数化简),可以减少不必要的访存,也更有利于生成高效算子。 图3.58代数化简对边界划分的影响 图3.59调整计算顺序 另外还可以对计算逻辑进行计算量上的优化。例如,可以在不影响计算逻辑结果正确性的前提下,调整计算顺序。如图3.59所示,通过调整铺展(Tile)的顺序,可以让求积(Square)的计算量减少,达到加速效果。 3.5.2MindSpore上的图算融合 图算融合在MindSpore上的整体架构如图3.60所示。主要分为前后端两部分,前端(Frontend)主要负责图层上的处理,后端(Backend)则负责算子的生成。计算图传给图算融合前端,经过系列变化后,得到算子的计算逻辑信息(图3.60中的Operator Info),图算融合后端收到计算逻辑信息后生成对应硬件的执行算子。 图3.60图算融合在MindSpore上的整体架构 前端主要包含以下四大块。 (1) 扩展复合算子(CompositeOp Expandsion): 主要将目标复合算子的计算逻辑扩展成计算子图。 (2) 算子聚合(Op Aggregation): 将基础计算、已被扩展出来的子图等聚合起来。 (3) 高层优化(HighLevel Optimization): 包含冗余消除、代数化简、CSE等计算逻辑优化功能。 (4) 边界重建(Fusion Reconstruction): 主要是在已聚合算子中,重新划分高效边界,形成新的算子。 后端由AKG接管,主要包含以下三部分: (1) Poly调度优化(Polyhedral Schedule): 通过Polyhedral技术对算子进行自动调度优化; (2) 低层优化(LowLevel Optimization): 贴近低层的优化,应用了包含双缓冲(doublebuffer)、内存复用、流水(pipeline)排布在内的多种优化技术; (3) 代码生成(CodeGen): 根据系列优化后的结果,生成对应硬件的逻辑代码(如GPU上生成高质量cuda代码)。 在MindSpore上使能图算融合,可以为网络执行带来可观的性能收益。在BERTBASE中,图算融合的优化数据如表3.6所示。 表3.6MindSpore+图算融合和 NGCTensorFlow+XLA的吞吐率(sentence/s)对比 节点数 GPU数 BatchSize MindSpore w/o 图算融合 MindSpore w/ 图算融合 NGCTensorFlow w/o XLA NGCTensorFlow w/ XLA 1 1 32 135.0 475.7 167.0 342.9 1 1 64 183.6 572.3 200.8 460.4 1 8 32 944.4 3276.4 1244.9 2437.6 1 8 64 1333.1 4287.6 1555.4 3368.5 3.5.3小结 图算融合在计算图中通过融合和跨算子优化技术重建优质子图结构,并基于Polyhedral技术进行高性能融合算子生成,实现高效的图算联合优化,带来显著的性能收益。在后续演进中,图算融合一方面将继续增强极致优化能力,在扩大融合范围、增强跨算子优化能力、提高设备使用率等方面持续发力; 另一方面,将增强图算融合的泛化能力,在面向更多样的计算图结构时,可以稳定、正确地运行,获取可观的性能收益。 3.6推理图优化 推理是指根据模型的输入和算子间的拓扑关系,逐层计算每个算子的输出,最终得到整个模型输出的过程。 常见的推理图优化技术包括模型剪枝、模型量化、算子融合、常量折叠、算子重排、算子替换等技术。TensorFlow Lite、MNN、MindSpore Lite等推理框架均实现了模型量化、算子融合、常量折叠和算子替换的图优化技术。MindSpore Lite和TensorFlow Lite还实现了算子重排的图优化技术。 这些图优化技术中模型剪枝和模型量化对于模型大小和模型推理时延带来的提升非常明显,但通常需要数据集的支持,而且可能丢失模型精度; 而算子融合、常量折叠、算子替换和算子重排等技术不依赖于数据集,不影响模型精度,通用性更强,应用更加便捷和广泛。本节主要介绍算子融合技术、算子替换技术、常量折叠技术和算子重排技术,并介绍相关技术在MindSpore上的实现和表现。 3.6.1算子融合 算子融合,顾名思义就是将深度神经网络模型中的多个算子,按照一定的规则替换成一个新的算子。后面将网络结果中需要被融合的多个算子及其拓扑关系称为融合模式(Pattern)。算子融合可以减少模型前向推理时的计算量,减少推理时的访存开销,降低推理时的时延和功耗,有些算子融合模式还能降低模型文件的大小。 算子融合带来的性能上的收益主要来自两个方面。一是通过融合,充分利用寄存器和缓存的空间,减少两次计算之间数据从内存中存储读取(storeload)的耗时。如图3.61所示,融合后,前一次计算的结果可以先暂存在CPU的寄存器(Register)或者缓存(Cache)中,下一次计算直接从寄存器或者缓存中读取,寄存器读取不到且缓存未命中(cache miss)才从内存中读取数据。二是通过融合,可以将一些计算量提前完成,避免了前向推理时的冗余计算或者循环冗余计算。 1. 计算公式融合 计算公式融合是指根据两个算子计算公式之间的链式调用关系,通过合并同类项、提取公因式等数学方法,将两个公式合并成一个公式,并且将合并得到的公式映射到某一种算子上,即完成多个算子的融合,并且在合并的过程中可以将一些参数的加减乘除提前计算完成。计算公式融合的优化效果很明显,因为该方案不仅避免了两个算子之间的数据访存过程,还将一些计算量提前在图优化中进行。 如图3.62所示,以Convolution算子和Batchnorm算子的融合为例,分解计算公式融合的过程,揭示计算公式融合的原理,并分析该图优化手段带来的收益。 图3.61数据存储和访问 图3.62Convolution算子与Batchnorm算子融合 在CV领域的模型中,为了保证每层算子的输入保持基本相同的分布从而加快收敛速度,Convolution算子之后经常会跟一个Batchnorm算子。针对这种融合模式,首先分析Convolution算子和Batchnorm算子的计算公式,卷积的计算公式为 Y=WX+B (3.40) 式中,W为Convolution的卷积核(weight),B为Convolution的偏差(bias),X为Convolution的输入,Y为Convolution的输出。 Batchnorm的计算公式为 Y=γ(X-μB)/δ2B+ε+β (3.41) 式中,X为Batchnorm算子的输入,μB为X的均值(mean),δ2B为X的方差(variance),γ为对X的缩放因子(scale),β为对X的平移因子(bias),Y是Batchnorm算子的输出。 当两个算子前后相连时,将Convolution的输出作为Batchnorm算子的输入 XBN=WX+B (3.42) Y=γ(XBN-μB)δ2B+ε+β (3.43) 简化后可以得到要更新的公式 Y=γWXδ2B+ε+γ(B-μB)δ2B+ε+β (3.44) 在实际应用中,该公式中除了Convolution的输入,其他符号都是常量,从而该公式代表的含义其实是一个Convolution算子,新的Convolution算子的filter和bias分别是 Wnew=γWδ2B+ε (3.45) Bnew=γ(B-μB)δ2B+ε+β (3.46) 可见,经过一次Convolution和Batchnorm算子的融合,一方面计算量从原来的一个Convolution和一个Batchnorm两个算子的计算量减少到一个Convolution的计算量; 另一方面也省去了Convolution算子的计算结果从寄存器写入缓存或者内存,然后Batchnorm算子又从缓存或者内存中加载卷积的计算结果到寄存器这个过程的时间。计算量降低后,模型推理时功耗自然而然就降低了。同时,由于删去了Batchnorm算子,Batchnorm算子对应的权重也可以删去,从而模型的大小也会降低。 如图3.63所示,在MindSpore Lite上以单Convolution算子和Batchnorm算子模型、mobilenet v2、nasnet mobile、inception v3、inception v4模型做一个简单的测试,经过Convolution和Batchnorm融合后,根据模型中Convlution算子和Batchnorm算子的比重,性能有不同程度提升。 图3.63Convolution算子与Batchnorm算子融合性能测试 2. 其他融合策略 除了计算公式融合,业界还有一些常用的融合策略,如将加法算子或者减法算子作为bias融合到前一个算子中,将激活算子集成到前一个算子中,将一系列小算子融合成一个业界常见的高级算法算子等。这里以Matmul算子和Add算子的融合、Convolution算子和ReLU算子的融合及GRU算子的融合为例,分别分析上述三种融合策略。 1) Matmul算子和Add算子的融合 Matmul算子的功能是计算两个矩阵的乘积。在深度神经网络模型中,经常会遇到Matmul算子后面跟一个Add算子或BiasAdd算子,对于这种常见的拓扑,如果Matmul算子除了两个矩阵的输入,支持bias作为第三个输入,那么就可以将Add算子或BiasAdd算子融合到Matmul算子中,以减少访存的开销。Matmul算子与Add算子融合的方式如图3.64所示。在MindSpore Lite中测试,该融合Pattern可以带来2.97%的性能提升。 2) Convolution算子和ReLU算子的融合 在深度神经网络模型中,为了增加模型的表达能力,通常会在Convolution算子后面增加一个激活函数,来增加模型的非线性表达能力。对于这种常见的拓扑,在Convolution算子中增加一个名为activation_type的属性,来表达融合到Convolution算子中激活函数的类型。同时后端Convolution算子的实现需要支持activation_type这个属性。如图3.65所示,Convolution算子后面跟随一个ReLU算子,那么融合得到Convolution算子的activation_type属性改为ReLU。 图3.64Matmul算子与Add算子融合的方式 图3.65Convolution算子与ReLU算子融合 3) GRU算子的融合 在深度神经网络算法中,存在一些由一系列计算过程组成的算法,每种算法都可以被表达为一个算子,比如业界常用LSTM算子和GRU算子。但这类算子为了编译训练,比如为了方便训练时进行自动微分,往往会被表达成一系列的小算子。如果不对这类算子做融合,一方面模型规模很大,不利于调试,另一方面模型中存在较多的小算子,这些小算子计算量不大,但访存开销却不小。 这里以GRU算子为例,GRU在TensorFlow中会被表达为20多个小算子,各个小算子之间数据输入输出均有访存开销。经过MindSpore Lite转换工具融合成单个GRU算子(如图3.66所示),可以带来35%的性能提升。 图3.66MindSpore中的GRU算子 3.6.2算子替换 算子替换,顾名思义就是将模型中某些低效的算子替换为高效的算子。算子替换的原理是通过合并同类项、提取公因式等数学方法,将算子的计算公式加以简化,并将简化后的计算公式映射到某类算子上。算子替换可以达到降低计算量、降低模型大小的效果。 这里以Batchnorm算子替换成Scale算子为例,分解算子替换的过程,揭示算子替换的原理,并分析该图优化手段带来的收益,如图3.67所示。 图3.67Batchnorm算子替换成Scale算子 如3.6.1节所述,Batchnorm的计算公式是 Y=γXδ2B+ε-γμBδ2B+ε+β (3.47) 在实际应用中,该公式中除了Batchnorm的输入和输出,其他符号均为常量。从而该公式可以归纳为一个变换和一个平移,可以用一个Scale算子来表达,新的Scale算子的scale和bias分别为 Scalenew=γδ2B+ε (3.48) Biasnew=β-γμBδ2B+ε (3.49) 可见,经过一次Batchnorm算子的替换,计算量从原来的一个Batchnorm算子减少到一次乘法和一次加法,计算性能可以得到提升,模型推理的功耗得到降低。同时Scale算子的权重比Batchnorm算子的权重少一半,整个模型也可以变小。 在MindSpore Lite上以单Batchnorm网络、mobilenet v2、nasnet mobile、inception v3、inception v4模型做一个简单的测试,将Batchnorm算子替换为Scale算子后,根据模型中Batchnorm算子的比重,性能提升5.9%~10.3%不等,如图3.68所示。 图3.68Scale算子替换Batchnorm算子性能测试 3.6.3常量折叠 1. 技术原理 常量折叠是一种常见的编译器优化技术,在编译时通过对常量或者常量表达式提前计算来提高运行效率,如代码3.19所示。 代码3.19常量折叠示例 i = 320 * 200 * 32; 编译器通常会在代码的中间表示(Intermediate Representation,IR)中直接计算出320×200×32的值,而不需要在运行中再计算i的值。另外编译器通常也支持常量传播,即逐步传播常量表达式计算出来的结果用来简化代码,示例代码如代码3.20所示。 代码3.20常量传播示例 int x = 14; int y = 7 - x / 2; return y * (28 / x + 2); 由于x为常量,在编译中将x代入第二行代码即可计算出y,再把y代入第三行代码即可直接计算出返回值。 在计算图优化中也借鉴了这种思想,输入全都是常量的节点可以在图编译时计算出执行结果并直接输出给下一个节点以减少运行的计算量。如图3.69所示,sqrt算子的输入只有一个常量输入,可以直接出结果并代入div算子计算出1/sqrt(w)作为mul算子的常量输入。 图3.69常量折叠原理 2. 用MindSpore实现常量折叠 MindSpore通过执行多次单层常量折叠来实现常量传播。图3.70中,第一次迭代最外层的A0和B1节点都只有一个常量输入可以被折叠得到状态②; 再执行第二次迭代,此时B0只有一个常量输入可以被折叠,C0和C1都有非常量输入不能折叠,折叠B0得到最终状态③。 图3.70常量折叠整体流程 单层常量折叠的流程如图3.71所示,遍历计算图节点找到各个节点是否有可以折叠的输入,将该节点命名为A,其可折叠节点命名为B。找到后执行节点B并将计算结果传给节点A,最后删除节点B即完成了一次常量折叠。 图3.71单层常量折叠流程 MindSpore中常量折叠性能测试结果如图3.72所示。 图3.72常量折叠性能测试结果 3.6.4算子重排 算子重排是指将模型中算子的拓扑按照某些规则进行重新排布,在不影响模型的准确率的前提下,降低模型对计算量的需求。常用的算子重排优化技术有针对Slice算子、StrideSlice算子、Crop算子等对feature map做裁切的算子的前移,Reshape算子和Transpose算子的重排,BinaryOp算子和Reshape算子的重排等。这里以Slice算子的前移为例,分析算子重排的原理和约束条件。 考虑到Slice算子是从输入的feature map中裁取一部分作为输出,是否可以将这个裁切的过程前移,提前对feature map进行裁切,那么后续算子由于输入的feature map减小,其计算量也会相应地减少,这就是Slice算子前移的原理。但是Slice算子的前移有一定的约束条件,如图3.73左侧,对于一般的element wise的算子,总是可以将slice前移,比较复杂的情况是要考虑element wise类算子进行broadcast的情况。如图3.74右侧,对于Matmul算子,需要重新计算Slice算子的begin end和size属性,才能将Slice前移到Matmul之前。 图3.73将Slice算子前移到Mul算子和Matmul算子之前 图3.74Slice算子不能前移到Convolution算子之前 如图3.74所示,对于Convolution算子,由于Convolution操作相对灵活,不同的卷积核对于feature map的操作不尽相同。除了比较特殊的卷积核的Convolution操作,比如point wise卷积,可以比较方便地将Slice算子前移,其他情况下的Convolution很难将Slice算子前移,所以Slice算子对于Convolution算子的前移方案缺少通用性。 MindSpore Lite中实现了Slice算子前移的优化器,支持将Slice算子前移到Softmax算子、Reshape算子、Matmul算子、FullConnection算子、Transpose算子、Add算子、Sub算子或Mul算子之前。使用MindSpore Lite中的Slice算子前移优化器对图3.74中的模型进行优化,推理性能提升了8.5%。 3.6.5小结 5G、物联网技术的落地和繁荣,势必会带来对端边设备上AI能力的极大需求,而推理图优化和算子性能优化在一定程度上成为端边设备上部署AI能力的基础和关键。在端边的分布式设备上,往往计算资源和存储资源十分有限,而且程序的响应时间和功耗对于用户的体验有较大的影响。因而轻量化、高性能、低功耗的推理框架成为AI领域的一个热门方向。不同于算子性能优化技术,图优化技术不依赖于设备的硬件能力,在提高推理性能的同时还能降低模型推理的功耗和模型大小,推理图优化技术的重要性可见一斑。 经过前面分析可知,图优化技术可以大幅缩减模型的大小、计算量和功耗。MindSpore后续除了增强当前框架下的图优化技术,比如增加更多的融合Pattern,还会重点开拓更多的图优化技术,比如实现模型剪枝的技术、模型稀疏数据压缩技术等。除此之外,MindSpore后续还希望做一些尝试,比如实现一种更加通用和便捷的模型自动融合技术等。 3.7kernel优化 随着AI技术的普及,业内各深度学习框架都在持续更新,以期望能够提供更便捷的使用操作和更优越的性能。对于一个固定的深度学习网络,训练和推理性能的瓶颈就在kernel即算子的运行时间。针对算子实现,可以考虑的优化手段有以下三种。 (1) 深度学习编译器: 具有代表性的是TVM,它的优点是隔离了算子计算操作与不同后端硬件,使用者只需写一次计算原语,TVM就可以自动生成不同硬件的代码; 缺点是TVM自动生成代码的性能和数据量紧密相关,缺少通用性。 (2) 对接高效计算库: 有一些深度学习框架将自己的算子库独立开源,例如端侧场景CPU领域有TensorFlow的XNNPACK、Pytorch的QNNPACK等。因此,可以采用模型定义+自研框架+三方算子库的方式,实现模型推理与训练,但是会受限于第三方算子库的性能。 (3) 手动实现: 针对不同硬件的各类算子,选取合适的算法,充分发挥各硬件的计算能力。 下面将以端侧推理框架MindSpore Lite在移动设备上的部署为重点,介绍kernel优化的相关内容,主要包含两个方面: 一是与硬件加速处理相关的优化; 二是算法层面的优化。 3.7.1硬件优化 当前MindSpore Lite主要应用在手机上,硬件的优化可以从以下几个方面入手。 (1) 选取性能更优越的后端设备: AI模型在手机各种处理器上实现推理,从CPU、GPU到NPU等,各种后端提供了不同的算力,关键在于发挥处理器的极致算力,加速深度学习中大量的计算过程,提升模型的推理速度。综合来说,几种处理器在深度学习领域的性能从高到低为: NPU>GPU>CPU。 (2) 异构调度: 框架也可以根据当前终端的硬件配置采用异构调度的方式,将模型整图进行切分,切分的子图通过调度器按后端设备的算力分发,充分利用该终端上设备的计算资源,减少推理时间,实现性能优化。 (3) 在kernel层实现多线程: 目前市面上的手机主要是CPU多核处理器,在执行网络推理时,算子与算子间存在前后依赖关系,但是同一个算子的大量数据之间可能互相无影响,在算子运行时(Runtime),可以采用多线程切分无影响的数据,分发到不同的CPU处理器上实现性能的数倍提升。 以上描述都是宏观角度的硬件优化方法,本节主要介绍实现算子运行时和ARM CPU处理器相关的优化手段。 1. 汇编语言 开发者使用的C++、Java等高级编程语言会通过编译器输出为机器指令码序列,而高级编程语言能做的事通常受编译器所限,汇编语言是靠近机器的语言,可以一对一实现任何指令码序列,编写的程序存储空间占用少、执行速度快、效率优于高级编程语言。 但是,汇编语言也存在难写难读、程序烦琐等缺点,所以在实际应用中,最好是程序的大部分用高级语言编写,运行性能要求很高的部分用汇编语言来编写,通过混合编程实现优势互补。深度学习中涉及了大量的计算,使用汇编语言能够给模型训练和推理性能带来数十到数百倍量级的提升。 2. NEON指令 最初的ARM指令集都是针对单个数据计算的,随着技术发展,逐步更新加入了并行指令。NEON寄存器和NEON指令就是专门针对大规模并行运算而设计的。 以ARMv8系列处理器为例,该CPU上有32个NEON 寄存器v0~v31,如图3.75所示,NEON寄存器v0可存放128bit的数据,即4个float32,8个float16,16个int8等。 图3.75ARMv8处理器NEON寄存器v0的结构 针对该处理器,可以采用SIMD(Single Instruction Multiple Data,单指令多数据)提升数据存取计算的速度。相比于单数据操作指令,NEON指令可以一次性操作NEON寄存器的多个数据。NEON指令较多,下面主要介绍在深度学习业务场景下一些常见的计算指令。 (1) float32指令: 对于浮点数,常见的运算的指令有fmul、fadd、fmla等。以fmla指令为例,该指令的用法为fmla v0.4s、v1.4s、v2.4s。fmla指令计算功能如图3.76所示,用于将v1和v2两个寄存器中相对应的float值相乘累加到v0的值上。 图3.76fmla指令计算功能 (2) float16指令: kernel优化的另一个点是用float16代替float32实现计算,这样在保证精度基本不下降的前提下,可以一次操作8个float16数据,理论性能提升一倍,指令通常与float32一致,寄存器的使用有差别,例如乘加指令为fmla v0.8h、v1.8h、v2.8h。 (3) int8指令: 整网性能优化的另一个方法是模型量化(或称模型压缩)。目前端侧支持的量化方式包含权重量化、量化感知训练和训练后量化。其中,量化感知训练和训练后量化保存得到的模型在实际推理时主要采用整型int8数据计算,在保证精度牺牲不多的情况下,提升推理的性能。常见的int8指令以sdot为例,用法为sdot v0.4s、v1.16b、v2.16b,如图3.77所示,一次操作可完成4组计算,每组计算实现 4对int8数据相乘求和累加到一个int32上,相较float指令大大提升了计算速度。 图3.77sdot指令计算功能 图3.78计算机存储设备 3. 汇编语言优化 对于已知功能的汇编语言程序来说,计算类指令通常是固定的,性能的瓶颈就在非计算指令上。我们主要从两个目标入手提升汇编语言的性能: 提高缓存命中率,优化指令流水。 如图3.78所示,计算机各存储设备类似于一个金字塔结构,最顶层空间最小,但是速度最快,最底层速度最慢,但是空间最大。L1~L3统称为cache(高速缓冲存储器),CPU访问数据时,会首先访问位于CPU内部的cache,没找到再访问CPU之外的主存,此时引入了缓存命中率概念来描述在cache中完成数据存取的占比。要想提升程序的性能,缓存命中率要尽可能高。 如图3.79所示,左图是一个循环乘加操作的指令流水,首先从内存中ldr加载数据到寄存器,fmla执行乘加计算,str指令将寄存器的数据写入内存,3次循环需要9次流水。右图是对该流水的重排优化,CPU中有多个NEON寄存器,因此在fmla指令执行的同时,可以向其他寄存器加载数据,将流水次数优化至5次。 图3.79循环乘加操作的指令流水 下面简单列举一些优化手段。 (1) 循环展开: 尽可能使用更多的寄存器,以代码体积换性能。 (2) 指令重排: 打乱不同执行单元的指令以提高流水线的利用率,提前有延迟的指令以减轻延迟,减少指令前后的数据依赖等。 (3) 寄存器分块: 合理分块NEON寄存器,减少寄存器空闲,增加寄存器复用。 (4) 计算数据重排: 尽量保证读写指令内存连续,提高缓存命中率。 (5) 使用预取指令: 将要使用到的数据从主存提前载入缓存,减少访问延迟。 本书只对汇编相关知识做了简单的介绍,要熟练使用汇编语言需要更加深入的学习。实际使用时,汇编语言性能调优是一个检验型的工作,需要多次尝试尽可能接近理论最优性能。 3.7.2算法优化 大部分的AI模型都包含卷积算子,且在推理过程中,卷积类算子计算大,占到了整网90%甚至更多的时间,本节主要介绍卷积算子算法方面的优化手段。 1. Img2col Img2col: 将一张图片根据卷积核的大小按对应位置展开成二维的矩阵。以一个简单的卷积为例,如图3.80所示,是一个3×3的输入d和2×2的卷积核g进行卷积运算。计算时用2×2的滑窗不断从输入矩阵中截取出数据,与卷积核的对应位置数据相乘累加得到输出。使用传统的for循环进行卷积操作,遇到数据量大的情况会非常费时,效率十分低下。针对矩阵运算的速度问题已经有很多研究, 图3.80输入3×3卷积核2×2示意图 如果能把卷积操作转化为矩阵操作,可以明显加快卷积速度,Img2col就是因此提出的。如式(3.50)所示,输入矩阵D的行向量对应一个输出点相关的所有输入,列为输出点的个数; 权重矩阵G的列向量对应一个输出点相关的所有权重值,列为1,这样就可以把卷积操作转化为矩阵乘法的操作。 D=d00d01d10d11 d01d02d11d12 d10d11d20d21 d11d12d21d22 G=g00g01g10g11T Y=D×G (3.50) 在常见的神经网络中,卷积的输入通常都是四维的,默认采用的数据排布方式为NHWC,如图3.81所示是一个通用卷积示意图。输入维度为(1,IH,IW,IC),卷积核维度为(OC,KH,KW,IC),输出维度为(1,OH,OW,OC)。 图3.81通用卷积示意图 对卷积的Img2col规则如下。如图3.82所示,对该输入做重排,得到的矩阵见右侧,行数对应输出的个数OH*OW; 每个行向量里,先排列计算一个输出点所需要输入上第一个通道的KH*KW个数据,再按次序排列之后的通道,直到通道IC。 图3.82输入Img2col的矩阵 如图3.83所示,对权重数据做重排。将1个卷积核展开为权重矩阵的一列,因此共有OC列,每个列向量上先排列第一个输入通道上KH*KW的数据,再依次排列后面的通道 图3.83卷积核Img2col的矩阵 直到IC。通过重排,卷积的计算就可以转换为两个矩阵相乘的求解。 2. GEMM 对两个矩阵相乘操作,从3.7.1节可以明显看到,如果输入矩阵在内存中行连续,卷积核矩阵在内存中列连续,那么在遍历访问数据时可以大大提升缓存命中率,这就是最简单的GEMM(General Matrix Multiplication,通用矩阵乘法)。 计算输出矩阵C=A×B,A为M×K矩阵,B为K×N矩阵,要求解C的一行数据,需要对A的一行循环遍历整个B,由于缓存空间有限,B的数据会随着遍历在缓存中加载、删除、更新。当计算下一行时,A也到达下一行,B要从第一列重新开始。在这个过程中,重复地从缓存中添加和删除B相同的数据。从这个优化点入手可以继续优化矩阵乘法的性能。 如图3.84所示,将输出拆分成m×n的小块,以m=4,n=4为例,计算C矩阵的这个块需要m行左矩阵数据和n列右矩阵数据。 图3.84GEMM分块示意图 采用传统的矩阵乘法,对应的伪码见代码3.21,通过三层循环来实现。采用分块计算,对应的伪码见代码3.22。 代码3.21矩阵乘法传统计算 for (int i = 0; i < M; i++) { for (int j = 0; j < N; j++) { for (int k = 0; k < K; k++) { C[i][j] += A[i][k] * B[k][j]; } } } 代码3.22矩阵乘法分块计算 for (int i = 0; i < M; i += m) { for (int j = 0; j < N; j += n) { for (int p = 0; p < m; p++) { for (int q = 0; q < n; q++) { C[i + p][j + q] = 0; } } for (int k = 0; k < K; k++) { for (int p = 0; p < m; p++) { for (int q = 0; q < n; q++) { C[i + p][j + q] += A[i + p][k] * B[k][j + q]; } } } } } 在K、m、n的内部循环里,尽量保证所要用到的A、B矩阵的数据能加载到缓存里,那么内部循环的数据加载到寄存器的延迟就会大大缩短,从而提升性能。要对C分块计算,就要保证每次分块计算需要的A和B在内存中是连续的,因此需要对A、B矩阵的数据进行重排。通过这种对内存重新组织的方式可以提升高速缓存命中率,从而优化性能。 分块策略的具体取值要结合硬件的高速缓存大小、寄存器数量等条件而定,对于ARMv8的硬件,在实现时取m=12,n=8; 针对ARMv7的硬件,在实现时取m=4,n=8。在实际实现时,Img2col和GEMM的组合可以同时进行,更加节省运行时间。 3. Winograd算法 卷积计算归根到底是矩阵乘法,两个二维矩阵相乘的时间复杂度是O(n3)。1969 年,Volker Strassen提出基于分治的Strassen算法,将矩阵乘法的复杂度降到了O(nlog7) 即O(n2.81)。在此之后,许多科学家致力于降低矩阵乘法时间复杂度的研究。在卷积算子的实现中,这里采用Winograd算法来降低矩阵乘法时间复杂度。 以一维卷积运算为例,记为F(m,r),其中m代表输出的个数,r为卷积核的个数。输入为d=d0d1d2d3,卷积核为g=g0g1g2T,该卷积计算可以写成矩阵形式,如式(3.51)所示,需要6次乘法和4次加法。 F2,3=d0d1d2 d1d2d3g0 g1 g2=y0 y1 (3.51) 可以观察到,卷积运算转换为矩阵乘法时输入矩阵中存在着重复元素d1和d2,因此卷积转换的矩阵乘法相对一般的矩阵乘法有了优化空间。可以通过计算中间变量m0~m3得到矩阵乘法的结果,见式(3.52)。 F2,3=d0d1d2 d1d2d3g0 g1 g2=m0+m1+m2 m1-m2+m3 (3.52) 式中,m0~m3分别如式(3.53)所示。 m0=d0-d2×g0m1=d1+d2×g0+g1+g22 m3=d1-d3×g2m2=d0-d2×g0-g1+g22 (3.53) 通过m0~m3间接计算r1、r2,需要的运算次数包括: 输入d的4次加法,输出m的4次乘法和4次加法。在推理阶段,权重的数值是常量,因此卷积核上的运算可以在图编译阶段计算,不计入在线的运行时间。所以总的运算次数为4次乘法和8次加法,与直接运算的6次乘法和4次加法相比,乘法次数减少,加法次数增加。在计算机中,乘法一般比加法慢,通过减少乘法次数,增加少量加法,可以实现加速。 计算过程写成矩阵形式,如式(3.54)所示,其中⊙为对应位置相乘,A、B、G都是常量矩阵。这里写成矩阵计算是为了表达清晰,实际使用时,按照式(3.53)手写展开的计算速度更快。 Y=AT[(Gg)⊙(BTd)] BT=10-10 0110 0-110 010-1 G=100 0.50.50.5 0.5-0.50.5 001 AT=1110 01-1-1 (3.54) 通常深度学习领域使用的都是二维卷积,将F(2,3)扩展到F(2×2,3×3),可以写成矩阵形式,如式(3.55)所示,各常量矩阵的值见式(3.54)。此时,Winograd算法的乘法次数为16,而直接卷积的乘法次数为36。 Y=AT[(GgGT)⊙(BTdB)]A (3.55) 整个计算过程在逻辑上可以分为4步,如图3.85所示。 图3.85整个计算过程 针对任意的输出大小,要使用F(2×2,3×3)的Winograd算法,将输出切分成2×2的块,找到对应的输入,按照上述的4个步骤,就可以求出对应的输出值。当然,Winograd算法并不局限于求解F(2×2,3×3),针对任意的F(m×m,r×r),都可以找到适当的常量矩阵A、B、G,通过间接计算的方式减少乘法次数。但是随着m、r的增大,输入、输出涉及的加法及常量权重的乘法次数都在增加,那么乘法次数带来的计算量下降会被加法和常量乘法所抵消。 在实际使用时,并不会对所有的卷积采用Winograd算法,Winograd算法只适用于stride和dilation均为1的卷积,并且需要考虑该方法的收益。对一个输入为(N,IH,IW,IC),卷积核为(K,KH,KW,IC),输出为(N,OH,OW,OC)的卷积来说,直接卷积的乘法计算次数见式(3.56)。 M=N×OH×OW×KH×KW×IC×OC (3.56) 使用Winograd算法F(m×m,r×r)的乘法计算次数见式(3.57)。 M′=N× OHm × OWm ×(m+r-1)2×IC×OC (3.57) 理论来说,若M>M′,Winograd算法的乘法次数就会降低,但是还需要考虑加法次数的增加及实际使用的硬件情况,这里不做赘述。 4. Indirect buffer 上述Img2col+GEMM的方法会带来内存的开销,针对卷积,一个输入会与多个卷积核进行运算,因此采用这种空间换时间的方法有明显的性能提升。但是这种方法不适合用在手表、音响等内存比较小的IoT设备上,并且对于depthwise算子收益不明显。下面介绍的Indirect buffer(间接缓冲区)可以克服这些缺点。 如图3.86所示是一个IC=OC=group的depthwise卷积,在进行深度卷积操作时,不同通道上的值不累加,因此不适合采用Img2col展开为矩阵的方法。 间接缓存中存放的是输入数据的地址,而非输入数据本身。针对输入数据,为OH×OW的每个点申请一个KH×KW的缓冲区,存放输入数据第一个通道的地址,例如对于第一个点output,保存的是0~8的地址。因此,该方法适用于NHWC格式即通道值连续的输入。 图3.86Indirect buffer算法示意图 缓冲区中保存的是地址信息,运算时通过在缓冲区行间的循环,根据该行缓冲区保存的指针找到对应的输入数据,再将输入数据与kernel相乘,从而间接地实现卷积操作,因此叫作Indirect buffer。 针对其他小算子,如Add、Mul等,在实现kernel的计算时,尽可能使用ARM CPU的汇编指令,性能都会比使用高级编程语言有显著的提升。算法优化与后端不是强相关的,可以结合硬件的具体情况将上述算法优化应用到各种硬件的kernel优化实现中,从而提升性能。 3.7.3小结 kernel优化在算子层面降低了MindSpore框架网络训练和推理的时延,提升了性能,降低了时间成本,是框架的核心竞争力之一。目前MindSpore框架中,针对Ascend、GPU、CPU等各种硬件,大都采用优化的计算方法,结合硬件的结构手写实现了算子的数据计算,加速计算过程,尽可能地减少算子的执行时间,发挥硬件的极致性能,从而实现kernel优化。 伴随着深度学习的发展和硬件的更新,kernel优化需要持续地进行。未来除了继续现有的优化之外,也可以探索新的方案: 算子内异构调度、自动生成高性能代码等。 3.8本章小结 当前深度神经网络技术应用场景越来越丰富,同时深度神经网络技术的各项指标要求也越来越高,如大规模神经网络模型训练时如何在有限的存储和算力条件下,使训练更快地收敛,使训练后模型精度更高; 在端边侧部署AI能力时如何在端边侧严格的功耗和存储限制下,更快地进行模型推理,等等。深度神经网络训练和推理优化技术一直是业界的核心话题。 MindSpore作为一个全场景的深度学习框架,在本章中,对MindSpore在云服务场景和端边场景,在CPU设备和GPU设备,在训练场景和推理场景等各种场景的优化手段进行了介绍,后续MindSpore会进一步深化现有的优化技术,不断探索前沿的优化算法,并针对不同的场景设计更强力的优化手段,打造业界高性能的全场景深度学习框架。