极智AI | 一文看懂 Google TPU 脉动阵列加速卷积计算原理

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

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

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

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