模型推理一文看懂 Google TPU 脉动阵列加速卷积计算原理

Posted 2021-08-05 极智视界

  本教程详细解释了 Google TPU 脉动阵列加速卷积计算原理。

  TPU 中计算卷积的方式和 GPU 不同，主要是依靠一种称为 “脉动阵列” 的硬件电路结构来实现的。脉动阵列的主体部分是一个二维的滑动阵列，其中每一个节点都是一个脉动计算单元，可以在一个时钟周期内完成一次乘加操作。脉动阵列中的行数和列数可以相等，也可以不相等，在每行每列的计算单元之间通过横向或纵向的数据通路来实现数据的向右和向下滑动传递。

  举例：如下图，特征图像 F 大小为 3 x 3，通道数为 2；卷积核 W 和 G 大小为 2 x 2，卷积核个数为 2。

  以这个例子来解释脉动阵列实现卷积加速的过程。在这个例子中，采用固定卷积核权重，横向和纵向脉动输入特征值和中间部分和的方式。如下图所示，首先将 2 个 2 通道 W 和 G 卷积核权重静态地存储到脉动阵列的计算单元中，其中对应同一个卷积核的 2 个通道权重被排列在同一列中。再将 2 个通道的输入特征图 F 排列展开，每一行都错开一个时钟周期，并准备依次输入到脉动阵列中。

  当要用 2 个通道输入特征图 F 和卷积核 W 及 G 进行卷积时，F 首先通过重排，通过左侧数值载入器输入到脉动阵列最左边的一列中。输入值从上到下，每一行进入脉动阵列的数据都会比上一行延迟一个时钟周期。在第一个时钟周期内，输入值 F00 进入 W00 乘加单元与 W00 权重进行乘加运算，产生单元结果 Y00，此为 Y00 的第一部分和。脉动阵列第 1 个时钟周期的计算状态如下图。

  在第二个时钟周期内，上一次 W00 乘加单元的部分和 Y00 被向下传递到 W01 乘加单元中，同时第二行输入值 F01 与 W01 相乘的单元结果与传递下来的 Y00 相加得到 Y00第二次部分和。而 F01 进入 W00 乘加单元中，进行乘加运算求出 Y01 的第一次部分和。与此同时 F00 向右滑动进入 G00 乘加单元中求出单元结果 Y10 的第一次部分和。脉动阵列第 2 个时钟周期的计算状态如下图。

  以此类推，输入的特征值沿着脉动阵列的行水平方向不断向右滑动开启不同卷积核各自的部分和。而对应于每一个卷积核的部分和沿着脉动阵列的列垂直方向不断向下滑动并和当前计算单元的结果相累加，从而在每一列的最下方的计算单元中得到最终累加完成的对应于该卷积核所有通道的卷积结果。

  这个例子中在第 8 个时钟周期会得到第一个卷积结果，如下图。

  而从第 9 个时钟周期开始，如下图。每次都可以得到两个卷积结果，直到所有的输入特征图像都被卷积完毕。脉动阵列通过横、竖两个方向的同时脉动传递，输入值呈阶梯状进入阵列，顺次产生每个卷积核的最终卷积结果。

  脉动阵列的特性决定数据必须按照事先安排好的顺序依次进入，所以每次要填满整个矩阵就需要一定的启动时间，而这段时间往往会造成硬件资源的浪费。启动时间通常可以通过脉动阵列的 “行数 + 列数 - 1” 计算得到。当度过启动时间后，整个脉动阵列进入满载状态，可以发挥出最大的吞吐率。

  在上面的例子中采用了固定卷积核权重，横向脉动输入特征值，纵向脉动部分和的方式来计算卷积结果。同理，也可以采用固定部分和，横向脉动输入特征值，纵向脉动卷积核权重的方式来计算得到同样的结果。事实上，可以任意选择脉动三个变量中的两个而固定另外一个的方式来实现卷积计算，具体的选择会依据各种实现方式的实际情况和应用需求来决定。

  需要指出的是，当计算的卷积核过大或者通道太多时，如果把所有通道的权重都串联在一列上，那么势必造成脉动阵列过于庞大或者畸形，不便于现实的电路设计。这时候往往采用分割计算、末端累加的方式来解决这个问题。系统会把多通道的权重数据分割成几个部分，每个部分都能够适合脉动阵列的大小。然后依次对每个部分的权重进行计算，计算完成的结果会临时存放在脉动阵列底部的一组加法累加器中。当换入另一组权重块后结果就会被累加，直到所有权重块都遍历完成，累加器累加所有的结果后输出最终值。

  TPU 中的脉动阵列仅仅完成了卷积的工作，而完成整个神经网络的计算还需要其他计算单元的配合。如下图，矢量计算单元通过输入数据端口接收来自于脉动阵列中的卷积计算的结果，通过激活单元中的非线性函数电路来产生激活值，激活函数可以根据网络需求来定制。在适量计算单元中，还可以通过归一化函数对激活值进行归一化，再经过池化单元就可以对激活值进行池化输出。这些操作都由队列模块进行控制。例如队列模块可以通过配置参数的端口指定激活函数、归一化函数或池化函数，以及处理步长等参数。矢量计算单元处理完成后将激活值发送到片上的统一缓冲区中暂存，作为下一层网络的输入。TPU 按照这种方法，一层一层地完成整个神经网络的计算。

  虽然脉动阵列主要用来加速神经网络的推理计算，但这个架构不光是只能处理卷积，它在进行通用矩阵计算时依然高效且强大，因此还可以用来完成除了卷积神经网络以外的其他一系列工作，比如全连接神经网络、线性回归、逻辑回归、分类（如K均值聚类）与视频编码和图像处理等。

扫描下方二维码即可关注我的微信公众号【极智视界】，获取更多AI经验分享，让我们用极致+极客的心态来迎接AI ！