张量处理单元(TPU)

2017年，谷歌对外宣布了一个张量处理单元(也就是TPU)—专门为机器学习而构建的定制应用专用集成电路（ASIC），一年后，TPU被转移到云端，并开放供商业使用。张量处理器（TPU）遵循CPU和GPU的路线，是Google定制开发的特定于应用程序的集成电路（ASIC），旨在加速机器学习的工作量，专为谷歌的TensorFlow框架设计的，实际是一个用于神经网络的符号数学库。

众所周知，TensorFlow的源码阅读起来是比较困难的，所以为了更好地理解张量的概念、TPU结构以及它的工作原理，接下来将尝试对TPU技术做一个简单的概述。

什么是张量？

在数学里面，张量是以多线性方式将几何向量、标量和其他类似对象映射到结果张量的几何对象。听起来可能很令人懵逼，简单地说，张量是一个广义矩阵，可以是一维矩阵（向量）、三维矩阵（数字的立方体）、甚至是一个0xD矩阵（单个数字）或更高维的结构，这很难想象，张量的维数称为它的秩。

张量的一个重要特征是它的动态特征：它生活在一个结构中，并与其他数学实体相互作用，如果以规则的方式变换结构中的其他实体，张量将遵循相关的变换规则，其数值将与其他实体一起显示，张量的这种“动态”性质是区分它与纯矩阵的关键。

在计算机科学中，张量是一个n维矩阵，类似于numpy数组（机器学习算法使用的基本数据结构），是TensorFlow中的基本操作单元，实际上就是使用了numpy。

张量和机器学习有什么关系？

ML运行所涉及的数学包括了以下这些对象的乘积和加法运算：

• scalar
• vector
• matrix

让我们来看一个简单的神经网络，TensorFlow Playground上的以下示例训练神经网络，根据训练数据集将数据点分类为蓝色或橙色：

训练后的神经网络用标签或预估值对数据进行分类，这个过程称为推理，为此，神经网络中的每个神经元都会进行以下计算：

• 将输入数据（x）乘以权重（w），以表示信号强度；
• 将结果相加以聚合神经元的状态；
• 使用激活函数（F）调节人工神经元的活动。

这里可以看到，一共有三个输入，只有两个神经元和一个完全连接的单层神经网络，你必须在权重和输入之间执行六次乘法，并将乘法相加成两组，每组三个，这个乘法和加法序列可以写成一个矩阵乘法，是运行一个经过训练的模型最需要计算的部分，每一个预测都需要许多步骤，将处理过的输入数据乘以一个权重矩阵和应用一个激活函数。

因此，从矩阵中乘片和取片需要大量的CPU时钟周期和内存，TPU的目的就是减轻这些特定的工作负荷。

TPU背后的逻辑

为了简化神经网络训练过程，TPU从多个方面解决了这个问题：

1.量子化运算

根据量子技术，在预设的最小值和最大值与8位整数之间的任意值的近似过程，TPU包含65536个8位整数乘法器；本质上，这种技术是将32位甚至16位数字的浮点计算压缩为8位整数，可以看到，连续的大数集（如实数）通过保持曲线转换为离散集（整数）：

量子化运算是第一个强大的工具，TPU使用，以减少神经网络预测的成本，而不是显著地损失的准确性。

2.注重数学推理

TPU的设计本身封装了神经网络计算的本质，TPU包括以下计算资源：

• 矩阵乘法单元（MXU）:65,536个矩阵运算的8位乘法和加法单位；
• 统一缓冲区(UB)：24MB SRAM寄存器；
• 激活单元（AU）：硬接线激活功能。

它们由十几个高级指令控制，这些指令集中于神经网络推理所需的主要数学运算，一个特殊的编译器和软件栈将所有来自于TensorFlow图的API调用转换为TPU指令。

3. 并行处理

TPU采用典型的RISC处理器为简单的计算提供指令，例如用处理单个或标量的操作乘以每个指令；就如前面讲到的，TPU包含一个矩阵乘法器单元，它被设计成一个矩阵而不是标量处理器，并在一个时钟周期内处理数十万个操作（相当于矩阵操作），使用这种矩阵处理器就像一次打印一整页文档，而不是挨个字符或逐行打印。

4.收缩阵列

TPU的核心是一个新型的矩阵乘法单元(MXU)，叫做收缩阵列，在传统的体系结构中(如CPU或者GPU)，计算值存储在寄存器中，程序告诉算术逻辑单元（ALU）哪些寄存器需要读取，要执行的操作（如加法、乘法或逻辑与）以及将结果放入的寄存器，这些操作的序列组成了程序，在MXU中，矩阵乘法多次重复使用输入以产生最终输出。一个值被读取一次，但用于许多不同的操作，而不将其存储回寄存器，ALU只执行固定模式中的乘法和加法，并且电线连接相邻的ALU，这使得它们既短又节能。

为了理解这个设计，可以想想心脏泵血——就像数据在芯片中以波的形式流动。

云化

TPU最令人兴奋的部分应该是谷歌已经在谷歌云平台上提供了这些芯片作为服务，并开源了它的技术，云TPU允许使用TensorFlow在谷歌的TPU加速器硬件上运行机器学习工作流，然后将TPU工作负载分发到TPU云服务器，如下所示：

云TPU旨在帮助研究人员、开发人员和企业构建能够利用CPU、GPU和TPU的TensorFlow计算集群，高级TensorFlow API可以帮助你在云TPU硬件上运行模型。

程序设计模型

由于PCIe总线的速度，与计算速度相比，云TPU和主机内存之间的数据传输速度较慢。由于主机和云TPU之间经常发生“乒乓”现象，云TPU大部分时间处于空闲状态，等待数据通过PCIe总线到达。为了提高效率，云TPU的编程模型将大部分培训转移到TPU，这样的话：

• 所有模型参数都保存在同一片高带宽存储器中；
• 许多培训步骤都是在一个循环中执行的，用于分摊在云TPU上启动计算的成本；
• TensorFlow服务器在“馈入”到云TPU硬件之前获取并预处理数据；
• 云TPU上的核心同步执行存储在各自HBM中的相同程序。

价格

根据谷歌的定价信息，每台TPU的成本为4.50美元，网站上显示的三种TPU型号的价格没有差异，由于TPU具有不同的内存大小和时钟速度，这有点让人懵逼。

CPU/GPU/TPU对比

CPU（中央处理单元）和GPU（图形处理单元）计算大多数ML问题时会消耗大量时间，一些专家认为，TPU本质上是Nvidia的GPU转移到了云上，由于这些技术非常相似，我们提供了一张对比CPU、GPU和TPU芯片的表格，以显示它们的差异以及它们的应用范围：

TPU优势

在提高计算效率和速度方面，使用TPU具有许多优势，包括以下好处：

• 线性代数计算的加速性能，在机器学习应用中被大量使用。
• 在训练大型、复杂的神经网络模型时，最大限度地缩短了训练的时间：以前训练几周的模型在TPU上可以在数小时内收敛。
• 通过其TPU服务器跨不同机器的可扩展操作。

TPU局限性

要理解的一件重要的事情是，TPU被专门优化以执行快速、庞大的矩阵乘法，在不受矩阵乘法支配的工作负载中，云TPU可能比其他平台表现更好，例如：

• 线性代数程序，需要频繁的分支或由代数控制的元素；
• 以稀疏方式访问内存的工作负载；
• 需要高精度算法的工作负载；

神经网络工作负载，包含用C++编写的定制的TysFoRoad操作，具体地说，是主训练循环主体中的自定义操作。

TensorFlow和TPU的结合可以在医学、图像处理和ML领域创造奇迹。当训练模型的时间和计算时间从几周缩短到几小时时，机器学习变得更具竞争力，并可供更多的人使用，此外，将TPU作为云上服务提供，允许用户开始构建他们的模型，从而无需进行前期较大的资本投资，这意味着研究人员、工程师、小企业甚至学生可以轻松地创建机器学习项目。