TensorFlow: A System for Large-Scale Machine Learning
1.Abstract
TensorFlow是在大规模和异构环境中运行的机器学习系统。 TensorFlow使用数据流图来表示计算,共享状态以及使该状态发生变异的操作。 它在群集中的许多计算机之间以及一台计算机中的多个计算设备之间映射数据流图的节点,这些计算设备包括多核CPU,通用GPU和称为Tensor处理单元(TPU)的定制设计ASIC。 这种体系结构为应用程序开发人员提供了灵活性:而在以前的“参数服务器”设计中,共享状态的管理已内置到系统中,而TensorFlow使开发人员能够进行新颖的优化和训练算法。 TensorFlow支持各种应用程序,重点是对深度神经网络的训练和推理。 一些Google服务在生产中使用TensorFlow,我们已经发布了itasan开源项目,并且已经广泛用于机器学习研究。在本文中,我们描述了TensorFlow数据流模型,并展示了TensorFlow在多种实际应用中实现的出色性能。
2. Background & motivation
2.1前身系统 DistBelief
DistBelief使用参数服务器架构,在这里我们批评它的局限性。
在参数服务器体系结构中,作业包括两个不相交的过程集:(stateless worker processes)无状态工作进程(在训练模型时执行大量计算)和(parameter server processes)有状态参数服务器进程(维护模型参数的当前版本)。 DistBelief的编程模型类似于Caffe的编程模型[38]:用户将神经网络定义为层的有向无环图,并以损失函数终止。 层layer是数学运算符的组合:例如,完全连接的层将其输入乘以权重矩阵,添加偏差矢量,然后应用非线性函数(例如sigmoid)。 损失函数是一个标量函数,它量化了预测值(对于给定的输入数据点)与基本事实之间的差异。 在全连接层中,权重矩阵和偏差向量是参数,学习算法将对其进行更新以最小化损失函数的值。 DistBelief使用DAG结构和层语义的知识,通过反向传播[63]为每个模型参数计算梯度。 因为许多算法中的参数更新是可交换的,并且对一致性的要求很弱[61],所以工作进程可以独立计算更新,并将“增量”更新写回每个参数服务器,从而将更新与其当前状态结合在一起。
尽管DistBelief使许多Google产品能够使用深度神经网络并构成了许多机器学习研究项目的基础,但我们很快就开始感受到它的局限性。 它的基于Python的脚本界面可用于组成预定义的层,足以满足具有简单要求的用户,但我们的高级用户寻求三种灵活性:
Defining new layers:为了提高效率,我们将DistBelief层实现为C ++类。 对于试图尝试新的层架构的机器学习研究人员来说,使用单独的,不太熟悉的编程语言来实现层是一个障碍。
Refining the training algorithms:许多神经网络都使用随机梯度下降(SGD)进行训练,该算法通过沿最大程度降低损失函数值的方向移动神经网络来迭代地细化网络参数。 通过更改更新规则[23,66],对SGD的一些改进可以加快收敛速度。 研究人员通常想尝试新的优化方法,但是在DistBelief中这样做涉及修改参数服务器的实现。而且,参数服务器的get()和put()接口并不是所有优化方法的理想选择:有时必须自动更新一组相关参数,在许多情况下,将计算卸载到参数服务器上会更有效 ,从而减少网络流量。
Defining new training algorithms:DistBelief worker遵循固定的执行模式:读取一批输入数据和当前参数值,计算损失函数(通过网络的正向传递),为每个参数计算梯度(向后传递),并编写梯度 回到参数服务器。 这种模式适用于训练简单的前馈神经网络,但不适用于更高级的模型,例如包含循环的递归神经网络[39];对抗网络,其中两个相关网络交替训练[26]; 强化学习模型,其中损失函数是由诸如视频游戏模拟器之类的单独系统中的某些代理计算的[54]。 此外,还有许多其他的机器学习算法,例如期望最大化,决策森林训练和潜在Dirichlet分配,它们与神经网络训练不尽相同,但也可以受益于通用的,优化的分布式运行时。
另外,我们在设计DistBelief时考虑了一个平台:一个大型的多核服务器分布式集群[20]。 当很明显,这种加速对于高效执行卷积内核至关重要时,我们能够添加对GPU加速的支持[44],但是DistBelief仍然是一个重量级系统,适合在庞大的数据集上训练深度神经网络,因此很难 缩小到其他环境。 特别是,许多用户希望在扩展相同代码以在更大的数据集上进行训练之前,在GPU驱动的工作站上本地磨练其模型。 在集群上训练模型后,下一步是将模型推入生产环境,这可能涉及将模型集成到在线服务中,或将其部署到移动设备上以供离线执行。 这些任务中的每一个都有一些共同的计算结构,但是我们的同事发现有必要使用或创建满足每个平台不同性能和资源要求的单独系统。 TensorFlow为所有这些环境提供了一个单一的编程模型和运行时系统。
2.2设计原则
我们设计的TensorFlow比DistBelief灵活得多,同时保留了满足Google生产机器学习工作负载需求的能力。 TensorFlow提供了一个简单的基于数据流的编程抽象,允许用户在分布式集群,本地工作站,移动设备和定制设计的加速器上部署应用程序。 一个高级脚本界面(图1)包装了数据流图的构造,并使用户能够在不修改核心系统的情况下尝试使用不同的模型架构和优化算法。 在本小节中,我们简要介绍TensorFlow的核心设计原则:
Dataflow graphs of primitive operators:TensorFlow和DistBelief都为其模型使用数据流表示形式,但最明显的区别是DistBelief模型包含相对较少的复杂“层”,而相应的TensorFlow模型将单个数学运算符(例如矩阵乘法,卷积等)表示为节点 在数据流图中。 这种方法使用户可以更轻松地使用高级脚本界面来编写新颖的图层。许多优化算法要求每一层都具有定义的渐变,并且通过简单的运算符构建层可以轻松自动区分这些模型(第4.1节)。 除了功能运算符之外,我们还表示可变状态以及将其更新的操作作为数据流图中的节点,从而可以尝试使用不同的更新规则。
Deferred execution:典型的TensorFlow应用程序具有两个不同的阶段:第一个阶段将程序(例如,要训练的神经网络和更新规则)定义为符号数据流图,并带有占位符表示输入数据和代表状态的变量; 第二阶段在一组可用设备上执行程序的优化版本。 通过将执行推迟到整个程序可用之前,TensorFlow可以通过使用有关计算的全局信息来优化执行阶段。 例如,TensorFlow通过使用图形的依赖关系结构向GPU发布一系列内核,而无需等待中间结果,从而实现了GPU的高利用率。 尽管此设计选择使执行效率更高,但我们不得不将更复杂的功能(例如动态控制流(第3.4节))推入数据流图中,以便使用这些功能的模型可以进行相同的优化。
Common abstraction for heterogeneous accelerators:除了多核CPU和GPU等通用设备以外,用于深度学习的专用加速器还可以显着提高性能并节省功耗。 在Google,我们的同事建立了专门用于机器学习的Tensor处理单元(TPU); 与替代的最新技术相比,TPU的每瓦性能提高了一个数量级[40]。 为了在TensorFlow中支持这些加速器,我们定义了设备的通用抽象。 至少,设备必须实现以下方法:(i)发出要执行的内核,(ii)为输入和输出分配内存,以及(iii)往返于主机内存的缓冲区。 每个运算符(例如矩阵乘法)可以具有针对不同设备的多种专门实现。 结果,同一程序可以轻松地将GPU,TPU或移动CPU定位为训练,服务和离线推理所需的目标。
TensorFlow使用原始值的张量作为所有设备都可以理解的通用交换格式。 在最低级别上,TensorFlow中的所有张量都是密集的; 稀疏张量可以用密集张量表示(第3.1节)。 该决定可确保系统的最低级别具有用于内存分配和序列化的简单实现,从而减少了框架开销。 张量还支持内存管理和通信的其他优化,例如RDMA和GPU到GPU的直接传输。
3 TensorFlow execution model
TensorFlow使用单个数据流图来表示机器学习算法中的所有计算和状态,包括各个数学运算,参数及其更新规则以及输入预处理(图2)。 数据流图显式表达了子计算之间的通信,因此使并行执行独立的计算以及跨多个设备划分计算变得容易。 TensorFlow与批处理数据流系统(§2.3)在两个方面有所不同:
a. 该模型支持对整个图的重叠子图执行多个并发执行。
b. 各个顶点可能具有可变状态,可以在图形的不同执行之间共享。
参数服务器体系结构中的主要观察结果是,可变状态对于训练非常大的模型至关重要,因为可以对非常大的参数进行就地更新,并将这些更新尽快传播到并行训练步骤中。具有可变状态的数据流使TensorFlow可以模拟参数服务器的功能,而且具有更大的灵活性,因为可以在承载共享模型参数的机器上执行任意数据流子图。我们的用户已经能够尝试使用不同的优化算法,一致性方案和并行化策略。
3.1 Dataflow graph elements
在TensorFlow图中,每个顶点代表局部计算的单位,每个边代表来自顶点的输出或输入。 我们将顶点处的计算称为运算operations,将沿边缘流动的值称为张量tensors。在本小节中,我们描述操作和张量的常见类型。
Tensors:在TensorFlow中,我们将所有数据建模为张量(n维数组),其元素具有少量原始类型之一,例如int32,float32或(其中字符串可以表示任意二进制数据)。 张量自然表示许多机器学习算法中常见数学运算的输入和结果:例如,矩阵乘法采用两个2-D张量并产生一个2-D张量; 批处理二维卷积需要两个4-D张量并产生另一个4-D张量。
由于我们在2.2小节中讨论的原因,所有TensorFlow张量在最低级别都是密集的。 TensorFlow提供了两种表示稀疏数据的替代方法:将数据编码为密集张量的可变长度字符串元素,或者使用密集张量的元组(例如,具有m个非零元素的n-D稀疏张量可以在坐标中表示为m×n的坐标矩阵和长度为m的值向量)。 张量的形状可以在其一个或多个尺寸上变化,这使得可以用不同数量的元素表示稀疏张量。
Operations:一个运算将m≥0张量作为输入,并产生n≥0张量作为输出。 运算符命名的“类型”(例如Const,MatMul或Assign),并且可能具有零个或多个确定其行为的编译时属性。 一个操作在编译时可以是多态的和可变的:它的属性决定了它的输入和输出的预期类型和一致性。例如,最简单的操作Const没有输入,只有一个输出。 它的值是一个编译时属性。 例如,AddN对相同元素类型的多个张量求和,并且具有定义其类型的属性T和整数属性N。
Stateful operations: variables:一个操作可以包含可变状态,该可变状态在每次执行时都会被读取和/或写入。Variable算子拥有一个可变缓冲区,该可变缓冲区可用于在训练模型时存储模型的共享参数。Variable没有输入,并产生一个引用句柄(reference handle),该句柄用作读取和写入缓冲区的类型化功能。 Read操作将参考柄r作为输入,并输出变量(State [r])的值作为密集张量。 其他操作会修改基础缓冲区:例如,AssignAdd使用参考句柄r和张量值x,并在执行时执行更新状态State '[r] ← State [r] + x。 随后的Read(r)运算产生值State'[r]。
Stateful operations: queues:TensorFlow包含几种队列实现,它们支持更高级的协调形式。 最简单的队列是FIFOQueue,它具有内部的张量队列,并允许按照先进先出的顺序进行并发访问。其他类型的队列以随机和优先级顺序使张量出队,以确保适当地采样输入数据。 像Variable一样,FIFOQueue操作生成一个引用句柄reference handle,该引用句柄可以由标准队列操作之一(例如Enqueue和Dequeue)使用。 这些操作将其输入推入队列的尾部,并分别弹出head元素并将其输出。 如果给定队列已满,则Enqueue将阻塞,如果给定队列为空,则Dequeue将阻塞。在输入预处理流水线中使用队列时,此阻塞会提供背压backpressure; 它还支持同步(第4.4节)。 队列和动态控制流(第3.4节)的组合还可以实现子图之间的流计算形式。
3.2 Partial and concurrent execution 局部和并发执行
TensorFlow使用数据流图来表示特定应用程序中的所有可能的计算。 用于执行图的API允许客户端以声明方式指定应执行的子图。 客户端选择零个或多个边以将输入张量馈入数据流,并选择一个或多个边以从数据流中获取输出张量。 然后,运行时将修剪该图以包含必要的操作集。 API的每次调用都称为一个步骤step,并且TensorFlow在同一图形上支持多个并发步骤。 有状态操作允许步骤共享数据并在必要时进行同步。
图2显示了一个典型的训练应用程序,其中有多个子图同时执行并通过共享变量variables和队列queues进行交互。 核心训练子图取决于一组模型参数和队列中的输入批次。 训练子图的许多并行步骤会根据不同的输入批次更新模型,以实现数据并行训练。 为了填充输入队列,并发的预处理步骤会转换单个输入记录(例如,解码图像和应用随机失真),一个单独的I/O子图从分布式文件系统中读取记录。 检查点子图会定期运行以提高容错能力(第4.3节)。
部分执行和并发执行使得TensorFlow具备了极高的灵活性。 通过队列添加可变状态和统一,可以在用户级代码中指定各种模型架构,这使高级用户可以进行实验而无需修改TensorFlow运行时的内部结构。默认情况下,一个TensorFlow子图的并发执行相对于彼此异步运行。这种异步性使得实现具有弱一致性要求的机器学习算法变得简单[61],其中包括许多神经网络训练算法[20]。 正如我们稍后讨论的那样,TensorFlow还提供了在训练期间同步工作人员所需的原语(第4.4节),这在某些学习任务(第6.3节)中产生了可喜的结果。
3.3 Distributed execution 分布式执行
数据流简化了分布式执行,因为它使子计算之间的通信变得明确。它可以将相同的TensorFlow程序部署到用于训练的GPU集群,用于服务的TPU集群以及用于移动推理的手机。
每个操作都驻留在特定设备上,例如特定task中的CPU或GPU。 设备负责为分配给它的每个操作执行内核。TensorFlow允许针对单个算子注册多个内核,并具有针对特定设备或数据类型的专门实现。 对于许多操作,例如按照元素的算子(Add,Sub等),我们可以使用不同的编译器为CPU和GPU编译内核实现。
TensorFlow运行时在图形上隐式或显式约束下对设备进行操作。放置算法为每个算子计算出一组可行的设备,计算必须共置的一组算子,并为每个共置组选择一个满意的设备。 它考虑了隐式的托管约束,因为每个有状态操作及其状态都必须放在相同的设备上。 另外,用户可以指定部分设备首选项,例如“特定任务中的任何设备”或“任何任务中的GPU”,并且运行时将遵守这些约束。 典型的训练应用程序将使用客户端编程结构来添加约束,例如,在一组“ PS”任务中分配参数。
因此,TensorFlow在如何将数据流图中的操作映射到设备上提供了极大的灵活性。 尽管简单的启发式方法为新手用户提供了足够的性能,同时专家用户可以通过手动放置操作以平衡多个任务和这些任务中的多个设备的计算,内存和网络需求来优化性能。 一个悬而未决的问题是TensorFlow如何在给定的一组设备上如何自动确定达到接近最佳性能的位置,从而使用户摆脱这种担忧。 即使没有这样的自动化,将放置指令与模型定义的其他方面分开也是值得的,这样,例如,在一个模型被训练之后修改放置就变得很简单了。
一旦在一个图中的算子被放置,并且部分子图已经被计算了一个步骤,TensorFlow就会将算子划分到每个设备的子图中。设备d的每个设备子图包含分配给d的所有操作,附加的Send和Recv操作替换设备边界上的边。一旦张量可用,Send就发送它的单个输入到一个指定的设备,同时使用一个约定的键rendezvous key来命名这个值。Recv只有一个输出,在产生该值之前,它会一直阻塞,直到指定键的值在本地可用为止。Send和Recv对几个设备类型对有专门的实现,在第5节中介绍其中一些。
我们优化了TensorFlow,以便在低延迟的情况下重复执行大的子图。一旦某个步骤的图被修剪、放置和分区,它的子图就被缓存到各自的设备中。客户端会话维护从步骤定义到缓存的子图的映射,因此可以通过向每个参与任务发送一条小消息来启动大型图上的分布式步骤。这个模型支持静态的、可重用的图形,但是它可以使用动态控制流支持动态计算,下一小节将对此进行描述。
3.4 Dynamic control flow 动态控制流
TensorFlow支持包含条件和迭代控制流的高级机器学习算法。例如,一个递归神经网络(recurrent neural network RNN)(如LSTM)可以从序列数据生成预测。谷歌的神经机器翻译系统使用TensorFlow来训练一个深层的LSTM,在许多翻译任务上实现了最先进的性能[73]。RNN的核心是一个递归关系, 其中序列元素i的输出是在整个序列上累积的某种状态的函数(图3)。 在这种情况下,动态控制流程可以在具有可变长度的序列上进行迭代,而无需将计算展开到最长序列的长度。
正如我们在2.2小节中讨论的那样,TensorFlow通过数据流图使用延迟执行将较大的工作量卸载到加速器上。 因此,为了实现RNN和其他高级算法,我们在数据流图中添加了条件(if语句)和迭代(while循环)编程结构。 我们使用这些原语来构建高阶结构,例如map(),fold()和scan()[2]。
为此,我们从经典的动态数据流体系结构中借用了Switch和Merge原语[4]。 Switch是一个解复用器demultiplexer:它接受一个数据输入和一个控制输入,并使用该控制输入选择其两个输出中的哪个应产生值。未接收的switch输出接收一个特殊的死值dead value,该死值递归地遍历图的其余部分,直到到达一个merge算子。 merge是一个多路复用器multiplexer:它将最多一个非死的输入non-dead input转发到其输出,或者如果两个输入均非死,则产生无效输出dead output。 条件算子使用Switch根据布尔张量的运行时值执行两个分支之一,并使用Merge合并分支的输出。 while循环更加复杂,并使用Enter,Exit和NextIteration运算符来确保循环的格式正确。
迭代的执行可以重叠,并且TensorFlow还可以在多个设备和进程之间划分条件分支和循环主体。 分区步骤添加了逻辑,以协调每个设备上每次迭代的开始和终止,并确定循环的终止。 正如我们将在4.1小节中看到的那样,TensorFlow还支持控制流结构的自动区分。 自动微分将用于计算梯度的子图添加到数据流图,该数据流TensorFlow跨潜在分布的设备进行分区,以并行计算梯度。
4 Extensibility case studies
通过为TensorFlow中的所有计算选择统一的表示形式,我们使用户可以尝试对DistBelief运行时进行硬编码的功能。 在本节中,我们讨论使用数据流原语和“用户级”代码构建的四个扩展。
4.1 Differentiation and optimization
许多学习算法使用SGD的某些变体来训练一组参数,这需要计算损耗函数相对于那些参数的梯度,然后根据这些梯度更新参数。TensorFlow包含一个用户级库,该库可区分损失函数的符号表达式并生成代表梯度的新符号表达式。 例如,给定一个由层和损失函数组成的神经网络,该库将自动导出反向传播代码。
微分算法执行广度优先搜索,以识别从目标操作(例如损失函数)到一组参数的所有反向路径,并对每个路径贡献的部分梯度求和。我们的用户经常对某些操作的梯度进行专门化处理,他们已经实现了批处理归一化[33]和梯度裁剪[60]之类的优化,以加快训练速度并使其更可靠。 通过将节点添加到图上以记录前向控制中的控制流决策,并在后向反向过程中反向重放这些决策,我们扩展了该算法,以区分条件和迭代子计算(第3.4节)。 在长序列上区分迭代计算可能导致大量中间状态累积在内存中,并且我们已经开发了用于在这些计算上管理有限的GPU内存的技术。
TensorFlow用户还可以尝试各种优化算法,这些算法可以在每个训练步骤中为参数计算新值。 SGD易于在参数服务器中实现:对于每个参数W,梯度∂L/∂W和学习率α,更新规则为W′←W −α×∂L/∂W。 参数服务器可以通过使用-=作为写操作。并在训练步骤之后将α×∂L/∂W写入每个W中来实现SGD。
但是,有许多更高级的优化方案很难用单个写入操作表示。 例如,动量算法根据每个参数在多次迭代中的梯度为每个参数累积“速度”,然后根据该累积计算参数更新; 已经提出了对该算法的许多改进[66]。 在DistBelief [20]中实现动量,需要修改参数服务器实现,以更改参数数据的表示形式,并在写操作中执行复杂的逻辑; 对于许多用户而言,这样的修改具有挑战性。 优化算法是积极研究的主题,研究人员已在TensorFlow之上实施了几种算法,包括Momentum,AdaGrad,AdaDelta,RMSProp,Adam和L-BFGS。 这些可以在TensorFlow中使用Variable运算和原始数学运算来构建,而无需修改底层系统,因此随着新算法的出现,很容易进行实验。
4.2 Training very large models
为了在高维数据上训练模型,例如文本语料库中的单词[7],通常使用分布式表示distributed representation,该表示将训练示例嵌入为跨多个神经元的活动模式,并且可以通过反向传播进行学习[30]。 例如,在一个语言模型中,一个训练示例可能是一个稀疏向量,其非零项对应于词汇表中单词的id,并且每个单词的分布表示将是一个低维向量[6]。
推论是通过将一批b稀疏向量与n×d嵌入矩阵相乘而开始的,其中n是词汇表中的单词数,而d是所需的维数,从而产生小得多的b×d密集矩阵表示形式; 为了进行训练,大多数优化算法仅修改通过稀疏乘法读取的嵌入矩阵的行。 在处理稀疏数据的TensorFlow模型中,n×d可以达到千兆字节的参数:例如,大型语言模型可能使用超过10的9次方个参数,词汇量为80万个单词,我们在文档模型[19]方面经验丰富 这些参数占用几个TB。这样的模型太大,无法在每次使用时复制给工作人员,甚至无法存储在单个主机的RAM中。
我们将TensorFlow图中的稀疏嵌入层实现为原始操作的组成部分。 图4展示了一个嵌入层的简化图,该图分为两个参数服务器任务。 该子图的核心操作是Gather,它从张量提取稀疏的行集,TensorFlow将该操作与其所操作的variable并置。动态分区(Part)操作将传入的索引划分为可变大小的张量,这些张量包含指定给每个分片的索引,而动态拼接(Stitch)操作将来自每个分片的部分结果重新组合为单个结果张量。 这些操作中的每一个都有一个对应的渐变,因此它支持自动微分(第4.1节),结果是一组稀疏更新操作,这些操作仅作用于最初从每个分片中收集的值。
通常,编写TensorFlow模型的用户不会手动构建图4所示的图。 相反,TensorFlow包含一些库,这些库公开切分参数的抽象,并根据期望的分布程度构建适当的基本操作图。
虽然可以在参数服务器中进行稀疏读取和更新[49],但TensorFlow可以灵活地将任意计算卸载到托管共享参数的设备上。 例如,分类模型通常使用softmax分类器,该分类器将最终输出乘以具有c列的权重矩阵,其中c是可能分类的数量; 对于语言模型,c是词汇表的大小,可以很大。 我们的用户已经尝试了几种方案来加速softmax计算。 第一种类似于Project Adam [14]中的优化,其中权重在多个任务之间分片,并且乘法和梯度计算与分片并置。使用采样的softmax [37]可以进行更有效的训练,该采样基于示例的真实类和一组随机采样的错误类执行稀疏乘法。 我们在第6.4节中比较了这两种方案的性能。
4.3 Fault tolerance
即使使用大量机器,训练模型也可能要花费数小时或数天的时间。我们经常需要使用非专用资源来训练模型,例如使用Borg集群管理器[71],这不能保证在训练过程中可以使用相同的资源。 因此,长时间运行的TensorFlow作业可能会遇到故障或抢占,并且我们需要某种形式的容错能力。 任务失败的可能性很小,以至于单个操作都需要容错,因此,像Spark的RDD [75]这样的机制会带来很大的开销,却收效甚微。 无需使对参数状态的每次写入都具有持久性,因为我们可以从输入数据重新计算任何更新,并且许多学习算法不需要强一致性[61]。
我们使用图形中的两个操作(图2)实现用户级检查点的容错能力:Save将一个或多个张量写入检查点文件,Restore从检查点文件中读取一个或多个张量。 我们的典型配置将任务中的每个变量连接到相同的Save操作(每个任务一个Save),以最大化分布式文件系统的I / O带宽。 Restore操作从文件中读取命名张量,而标准的Assign将恢复的值存储在其各自的变量中。 在训练期间,典型的客户端会定期运行所有save操作以产生新的检查点; 客户端启动时,它将尝试还原最新的检查点。
TensorFlow包含一个客户端库,用于构造适当的图形结构,并在必要时调用Save和Restore。这种行为是可定制的:用户可以对模型中的变量子集应用不同的策略,或者定制检查点保留方案。例如,许多用户保留自定义评估指标中得分最高的检查点。实现也可重用: 它可以用于模型微调和无监督的预训练,这是迁移学习的形式,其中模型的参数在一项任务上进行训练(例如,识别一般图像)用作其他任务(例如识别狗的品种)的起点。 将检查点和参数管理作为图形中的可编程操作可以使用户灵活地实现我们无法预料的其他方案。
检查点库不会尝试生成一致的检查点:如果训练和检查点同时执行,则检查点可能不包括训练步骤中的任何更新,全部更新或部分更新。 此行为与异步SGD的宽松保证兼容[20]。 一致的检查点需要额外的同步,以确保更新操作不会干扰检查点。 如果需要,可以在下一个小节中使用该方案在同步更新步骤之后获取一个检查点。
4.4 Synchronous replica coordination
SGD对异步具有鲁棒性,许多系统使用异步参数更新来训练深度神经网络,它们被认为具有可扩展性,因为它们在存在散乱的情况下仍保持高吞吐量。 吞吐量的增加是以在训练步骤中使用陈旧的参数值为代价的。 最近有人重新考虑了同步训练无法扩展的假设[10,18]。 由于GPU可以进行数百(而不是成千上万)台机器的训练,因此同步训练(就质量时间而言)可能比同一平台上的异步训练更快。
尽管我们最初将TensorFlow设计为用于异步训练,但我们已经开始尝试同步方法。 TensorFlow图使用户可以在训练模型时更改参数的读取和写入方式,并且我们实现了三种选择。 在异步情况下(图5(a)),每个worker都将在每个步骤开始时读取参数的当前值,并在最后将其梯度应用于(可能不同的)当前值:这种方法可确保高利用率,但每个步骤使用过时的参数值,使每个步骤的效率降低。 我们使用队列来实现同步版本以协调执行:阻塞队列充当障碍,以确保所有工作程序读取相同的参数值,每个变量队列从所有worker那里累积梯度更新,以便以原子方式应用它们。 简单的同步版本(图5(b))在应用它们之前会累积所有worker的更新,但是工作缓慢的workers会限制总体吞吐量。
为了减轻这些worker的负担,我们实现了备份工作程序backup workers(图5(c)),类似于MapReduce备份任务[21]。 MapReduce在检测到straggler之后启动备份任务,而我们的备份工作程序则是主动运行的,并且聚合会生成n个更新中的前m个。 我们利用了SGD在每个步骤中随机采样训练数据这一事实,因此每个worker都处理一个不同的随机批次,并且如果忽略特定批次则不是问题。
5 Implementation
TensorFlow运行时是一个跨平台的库。 图6说明了它的体系结构:C API将不同语言的用户级代码与核心运行时分开。
TensorFlow核心库是用C ++实现的,以提高可移植性和性能:它运行在多种操作系统上,包括Linux,Mac OS X,Windows,Android和iOS; x86和各种基于ARM的CPU体系结构; 以及NVIDIA的Kepler,Maxwell和Pascal GPU微体系结构。 该实现是开源的,我们已经接受了一些外部贡献,这些力量使TensorFlow可以在其他体系结构上运行。
分布式主服务器distributed master将用户请求转换为跨一组任务的执行。 给定一个图和步骤定义,它将对图进行修剪(第3.2节)和分区(第3.3节)以获得每个参与设备的子图,并缓存这些子图,以便可以在后续步骤中重新使用它们。 由于主服务器可以看到每个,因此它应用了标准优化,例如常见的子表达式消除和恒定折叠; 修剪是消除无效代码的一种形式。 然后,它跨一组任务协调优化子图的执行。
每个任务中的数据流执行程序dataflow executor都会处理来自主服务器的请求,并调度包含本地子图的内核的执行。 我们优化了数据流执行器,以较低的开销运行大型图形。 我们当前的实现每秒可以执行10,000个子图(第6.2节),这使大量副本可以进行快速,细粒度的训练步骤。数据流执行程序将内核调度到本地设备,并在可能的情况下并行运行内核,例如通过使用多个CPU内核或GPU流。
运行时包含200多种标准操作,包括数学,数组操作,控制流和状态管理操作。 许多操作内核是使用Eigen :: Tensor [36]实现的,该工具使用C ++模板为多核CPU和GPU生成高效的并行代码。 但是,我们自由地使用cuDNN [13]之类的库,在这些库中可以实现更有效的内核。 我们还实现了量化,从而可以在移动设备和高吞吐量数据中心应用程序等环境中更快地进行推理,并使用gemmlowp低级精度矩阵库[35]来加速量化计算。
我们专门针对每对源设备和目标设备类型进行Send和Recv操作。 本地CPU和GPU设备之间的传输使用cudaMemcpyAsync()API进行重叠计算和数据传输; 两个本地GPU之间的传输使用DMA缓解主机压力。 对于任务之间的传输,TensorFlow使用多种协议,包括基于TCP的gRPC和基于聚合以太网的RDMA。 我们也正在研究使用集体操作的GPU到GPU通信的优化[59]。
第4节描述了我们在C API之上完全用用户级代码实现的功能。 通常,用户编写标准操作以构建更高级别的抽象,例如神经网络层,优化算法(第4.1节)和分片嵌入计算(第4.2节)。 TensorFlow支持多种客户端语言,并且我们优先考虑Python和C ++,因为我们的内部用户最熟悉这些语言。 随着功能的确立,我们通常将其移植到C ++,以便用户可以从所有客户端语言访问优化的实现。
如果将子计算表示为操作的组成部分很困难或效率低下,则用户可以注册其他内核,这些内核提供用C ++编写的有效实现。 我们发现手动执行融合内核对于某些性能至关重要的操作(例如ReLU和Sigmoid激活函数及其相应的梯度)是有利可图的。 我们目前正在研究使用基于编译的方法进行自动内核融合。
除了核心运行时之外,我们的同事还构建了一些工具来帮助TensorFlow用户。 其中包括为生产中的推理提供服务的基础结构[27],使用户能够跟踪训练运行进度的可视化仪表板,可帮助用户理解模型中的连接的图形可视化工具以及用于跟踪模型执行情况的分布式探查器。 跨多个设备和任务的计算。 我们在扩展的白皮书中描述了这些工具[1]。