1.4.4 机器学习算法的新引擎——TPU和NPU
先从一个小故事讲起。
小故事:欺骗雷达的空袭
Y国为了实现其战略目标,调集了14架战斗机空袭K国的重要目标。为了让这14架战斗机躲过防空雷达的探测,Y国空军进行了针对性的训练,让14架战斗机组成密集的编队,在雷达看来,这似乎是一架民航客机,从而骗过了K国的防空系统。
K国的防空系统之所以没有正确识别Y国的战斗机编队,是因为当时的雷达并没有安装计算机识别系统,而是靠肉眼识别反射波。二战时期,日军偷袭珍珠港时,美国雷达站的士兵也将日军的第一攻击波机群误以为美军B17轰炸机,从而忽略了预警信息。
在AI时代,前面小故事中的欺骗战术将难以为继,这是因为,以图像模式匹配为代表的图像识别算法,能够使计算机高效地帮助人类识别图像,避免被常见的视觉误区干扰。
图像模式匹配的核心算法是矩阵乘法运算。打开任意一本高等代数或线性代数教科书,都可以很快理解矩阵乘法运算:
可以看出,对于N×N的矩阵运算,需要计算N2个元素,对每个元素都需要计算一次N维向量卷积,总共是O(N3)的时间复杂度。如果雷达回波图形被离散化为256×256的矩阵,而用于匹配的模式矩阵为16×16的矩阵,那么每次识别图像都需要计算(256-16)×(256-16)=57600次矩阵乘法,每次矩阵乘法都需要进行16×16×16=4096次乘加运算。也就是说,每次识别图像总共需要进行约2.36亿次乘加运算。如果雷达回波图形和模式矩阵的分辨率进一步提升,那么每次识别图像所需要的计算量也会随之进一步增加。
与此类似,深度神经网络训练算法所需的向量卷积运算量,也会随着深度神经网络的层数和节点数增加,呈立方律爆发式增长。
面对这一挑战,业界的不同厂商给出了不同的解决思路。Nvidia的解决思路是,在GPU的CUDA Core之外,增加专有的硬件计算单元Tensor Core,实现在一个周期内计算一定大小以内的矩阵乘加,比如在Nvidia Volt架构中集成的初代Tensor Core,就可以在一个指令周期内进行半精度浮点数构成的16×16矩阵的乘加运算。
另一种思路来自Google。Google认为,普通的GPGPU考虑了太多矩阵算法和向量算法以外的计算与实现,对于以深度神经网络和图像模式匹配为代表的AI算法,这些功能实属浪费。因此,Google研制了非常擅长计算矩阵和向量,而精简了其他计算功能的专精装备——TPU(Tensor Processing Unit,张量处理单元),外观如图1-7所示。
图1-7
“Tensor”这个词在这本书中经常出现,这里先介绍其概念。
名词与概念:Tensor
Tensor(张量)是一个数学名词,代表一个代数对象,可将其视为一个矩阵在阶数上的推广。当张量的阶为0时,张量退化为标量,也就是1个元素的向量;当张量的阶为1时,张量退化为向量,也就是1×n的矩阵;当张量的阶为2时,张量就是矩阵。张量的阶还可以为3、4、5……直到无穷大。
Google的TPU在名义上是张量专用的计算处理器,实际上只是简单且直接地在芯片内部集成了数万个计算单元,这种计算单元只能用于进行乘加运算。比如Google用于训练AlphaGo的下一代AlphaZero的Cloud TPU v2,就集成了128×128=16384个计算单元,可以并发地进行16384次乘加运算。
TPU在计算流程方面,也与CPU、GPU有较大的差异,其对比如图1-8所示。无论是CPU还是GPU,所有运算过程的中间结果都需要被保存到内存中。而TPU根本没有将中间结果保存到内存中,而是在执行完毕后直接将中间结果传递给下一步骤。
图1-8
对照图1-9所示的深度神经网络的结构设计图,我们就能够理解Google TPU这一设计的价值所在了。
图1-9
在深度神经网络的结构设计中,在输入层和输出层之间一般会有若干隐藏层,而各层的每个节点实际上就是对来自各上级节点的值及模型中各上级节点的权重做向量卷积运算,并将向量卷积运算结果传递给下一层的各个节点。
我们会发现,无论是通用的CPU还是传统的GPU,每层节点进行向量卷积运算的结果都必须被保存到RAM中,下一层节点将从RAM中读取这个值,而RAM的读取和写入时延一般为100ns左右的数量级,在这段时间内,GPU的每个CUDA Core都可以完成100次左右的运算。而Google在研发TPU时,基于深度神经网络算法的特征,在硬件设计层面实现了这一优化,在神经网络算法的执行效率(时间复杂度)上也实现了数量级的提升。
可惜的是,TPU是一个专供产品,Google以外的用户如果期望使用TPU进行大模型等深度神经网络算法的训练和推理,那么只能使用Google的云服务。