天奇:解构TensorFlow, 学习MXNet -- 新一代深度学习系统的核心思想

作者:陈天奇 ,本文首发于作者的个人微博

这是一个以标题党风格的科普文。深度学习工具潮流滚滚,各种工具层出不穷。也有各种文章从易用性,可移植性,灵活性和效率方面对于各个系统进行比较。这篇文章希望从系统设计上面来讲来回答这个讨论这个问题:如果想到从头设计一个TF一样的新一代深度学习系统,到底需要把握哪些要点。

计算单元:从layer abstraction到operator

大家熟悉的第一代深度学习系统,以cuda-convnet2和caffe为代表。这些系统主要的一大特点是提出了一个以深度学习计算层次layer为基本单元的计算单位。不同的layer可以包含权重,并且相互组合。因为有了可以被组合的概念,不同的层次可以被组合在一起,从而可以比较灵活地尝试各种网络结构。

包括Tensorflow在内的比较新的系统一般都采用了和layer稍微有一些区别的模式: operator。operator和layer最大的差别就是不再区分参数和数据,并且参数本身不再是layer的一个成员。这样的好处是可以很容易地实现如不同层之间共享参数的目标。并且可以使得某一层的参数本身也是由某个网络预测得到。当然和layer的模式相比,因为operator本身一般不带有内部状态,会对于一些实现造成一定困难。

更多资源:自动并行调度

新一代系统的一个必不可少的元素就是并行调度算法。 开始的深度学习系统都是针对一个GPU来进行设计。当多卡并行的需求出现之后,并行调度算法的需求就开始体现出来了。一开始的cuda-convnet2和cxxnet都是采取手工编写并行模式的方法。采用多线程进行显卡之间的通信。随着并行模式的多样性和多机多卡等复杂需求的出现,一个系统的调度算法开始变得必不可少。一个深度学习的调度算法需要具有这些特点:

  • 自动并行可以并行执行的程度
  • 同时进行参数拷贝,网络传播和显卡计算
  • 及时触发相关依赖的状态。

具一个简单的例子,在一般的数据并行的场景下,当最高层的梯度被计算完毕之后,我们可以继续反向传播计算更低层的梯度,而与此同时我们必须直接开始吧这个高层计算的梯度拷贝到其他显卡进行梯度的整合。当整合完毕之后又必须马上触发参数更新。这一系列的过程都是和梯度计算同时在不同的显卡下完成的。并行调度算法可以自动地分析计算之间的依赖关系,并且实现这一系列复杂的调度过程,把用户从并行编程的枷锁下解放出来,从而可以实现各种奇葩的并行模式。

更少损耗:基于图的优化

因为引入了计算图的概念。我们可以先像优化数据库系统一样对于计算流图进行各种优化。一个最常见的优化就是通过分析依赖关系来减少深度学习的内存损耗,利用更少的显卡资源来跑更深的网络,或者在资源有限的移动资源上面部署。这一点Tensorflow做的暂时并不好。MXNet在这个方面做了一系列领先的工作,相信未来这样的思想也会被吸收到各个其他工具中去。

灵活性和更快速的交互: TensorFlow欠缺的元素

基于计算图声明的方式虽然可以获得比较多的优化空间。但是这样的编程模式需要用户把所有的操作都声明在一张图里面。与之相对的,像torch, numpy这样直接通过把计算拆开一步一步命令声明的方式更加灵活。命令声明的方式可以允许我们在多个图之间相互交互,并且根据当前的情况运行不同的来决定到底要运行哪一部分。虽然Tensorflow这样的图计算框架也可以分步执行多种操作,但是每个图操作之间的时间损耗相对很大,也无法得到有效的并行。

一边是可以增加效率的声明式图编程,另外一边是更加灵活的过程式编程。如何把这两者有机地结合在一起是包括Tensorflow在内的所有深度学习系统需要思考,并行进行重新设计的问题。

整理这些架构思想?学习MXNet
不管你使用什么框架。如果大家对这些设计感兴趣,并且希望学习如何架构一个包含这些思想的深度学习系统,不妨学习一下MXNet 。 MXNet利用更加简洁的设计实现了这里介绍的所有思想。并且引入了包括内存优化和灵活交互调度等Tensorflow不具备的元素。 同时MX也具有更多的可扩展性,比较容易引入如Torch直接整合,分布式parameter server等各种新特性。

和一般深度学习系统不同的是。除了使用文档,MXNet本身的文档包含了对于模块设计,模块关系和设计算法的具体介绍,并且也有一些同学参与了中文版本的翻译。因为有了这些材料,基于mxnet可以更加容易地掌握所有这些核心的设计思路,并且更加轻便地打造和TF一样甚至超越它的的新一代深度学习系统。

链接: https://github.com/dmlc/mxnet

 

Report Story