TensorFlow XLA初步接触

Tensorflow越来越像一个编译器,把计算图编译为可执行代码。其中关键的部分就是XLA (Accelerated Linear Algebra)。我在实际使用中真切感受到了XLA带来的提升,希望对XLA能有更多的了解,因此花了点时间探索了一下。


XLA框架

关于XLA,Tensorflow给出了比较简略的说明。XLA主要是用来提升计算速度、节省内存(显存)等。XLA的输入语言称为“HLO (High Level Optimizer)”,HLO定义了整个计算图。随后,XLA对HLO进行一些机器无关、高层的优化,然后用LLVM等进行机器相关、底层的优化并生成代码。这个流程如下图所示。

  1. 描绘计算图HLO,这一步可以通过tf2xla、xla_client等完成
  2. 对HLO进行广义优化(机器无关),如CSELoop Fusion等编译常用优化策略
  3. 针对特定设备,对HLO进行优化
  4. LLVM等生成可执行代码
XLA流程,见https://www.tensorflow.org/xla/overview

Operation Semantics

我不是非常理解Operation Semantics是什么意思,有兴趣可以看看。我只知道HLO支持很多操作,其中比较容易接受的是Element-wise unary functions(包括abs、cos等)、Element-wise binary arithmetic operations(如相加、相乘等)……

跟这一部分比较相似的是TensorRT定义的操作,如IActivationLayerIElementWiseLayer等。其实神经网络很简单,靠这25个操作就能定义大部分网络。
有了这些定义之后,我们就可以描述一个网络即计算图。有了精确定义的计算图,就可以对其进行优化。


计算图优化

前面也提过,XLA首先进行计算图优化主要是跟机器无关的、高层的。如CSELoop Fusion等编译常用优化策略。下面出现的XLA计算图和及其优化的中间结果,可以通过设置环境变量来导出,然后转换一下即可。

CSE

我们先看看CSE(Common subexpression elimination)。我也把这出成了Byte Camp的题,由于我难以描述清楚题目,没被采用。其实这题很能打“我是搞深度学习的,为什么让我做这么多编程题”的脸。

假设我们在使用Tensorflow等编写神经网络时,为了使代码逻辑清晰,可能会写出如下运行时低效的计算:

在Tensorflow中可以表示为如下左图,其中 p1 / (p0 + (p3 - p4)) 计算了两次。XLA就能对此进行了优化,只需计算一次,计算流程被优化为如下右图形式。
   
我们可以通过简单的程序来完成这一过程,可以看到真实的Tensorflow代码才200行不到。这也是ICPC的一道题,有兴趣可以尝试一下。

Fusion

Fusion可能带来提升,有可能会降低效率。这跟计算和架构相关。但在神经网络和Nvidia的GPU架构下,很难出现效率降低。看一个简单的例子, np.sin(np.cos(a * b) + c) ,其中 a,b,c 都是矩阵。显然通过fusion,我们可以,

  1. 节省存储,提高cache利用率
  2. 减少kernel数

Fusion后的计算表示为CUDA代码,大概是:

可以看到本来要4个CUDA Kernel要完成的计算,Fusion之后1个就行了。这提升还是非常靠谱的,XLA也对此做了优化,如下图所示。显然可以带来计算效率的提升 (gpu额外开销比较大)。

Fusion后Fusion后


BERT XLA

前面的都是随便写的计算,现在可以看看BERT开启XLA后发生了什么。先放一张图。整个BERT的计算图太大了,放不下。这里是一层Transformer,不带训练的情况,其实也够看了(放大看)。

BERT,一层Transformer的计算流程

从上图,我们可以看到Layer norm、GELU都有很多细碎的操作,这如果没有优化会产生很多额外开销和中间结果,带来的后果就是显存占用高。而XLA将这些细碎操作都Fusion在一起了,形成了一个大的Kernel。开启XLA和FP16之后,训练效率是原来的4倍,直接起飞,可能这个加速比还不是理论极限。


XLA Client

还有一个比较有意思的是XLA Client。这非常硬核,我们可以直接将numpy代码转成在GPU上运行的代码,并且附带计算图优化,完成了下面几个项目大部分功能。

  1. https://github.com/andersbll/cudarray
  2. https://github.com/cupy/cupy
  3. https://devblogs.nvidia.com/numba-python-cuda-acceleration/
  4. https://github.com/dmlc/minpy

具体可以参考,

  1. https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/xla/python/BUILD#L228
  2. https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/tf2xla/python/xla.py
  3. https://github.com/google/jax

这里简单使用一下xla client(基于tf 1.13.1),由于这个功能还不稳定,这个代码随时跑不起来。

计算速度如下,

计算图我就不放了,这个页面太大了,已经很卡了。。。有兴趣可以自己输出计算图看看。

 

参考链接

  1. https://www.tensorflow.org/xla/overview
  2. http://pages.di.unipi.it/corradini/Didattica/PR2-B-14/OpSem.pdf
  3. https://www.cs.cmu.edu/~rjsimmon/15411-f15/lec/18-commonsub.pdf

Go接触

最近有时间也有必要接触了Go,也有很多原因

  1. 大三的时候就想学学Go,不过由于Python耽搁了
  2. Python的并发让我头很疼,一点也不优雅,而Go的天然并发早声名在外
  3. 业务需求
  4. 多学一门语言,提升自己的竞争力

Go的特点

学习Go几乎没有阻碍,跟着A Tour of Go一步步学下来就行。光学怎么写Go就太枯燥了。我顺便了解了一下Go,很多内容在wiki中已经说得很清楚。总结一下,

  1. Go是静态类型,与C/C++类似,不同于Python,因此性能肯定比基于解释器的Python高。Go有垃圾回收机制、天然并发,写起来就很舒服,代价是性能肯定没有C那么高,看到一句话很有意思:Informally, Go is 50% of C’s performance (which is still really good) for 0% of the headaches.
  2. Go这个名字取得很好,Go由Google成员开发,Go是围棋,也是gopher(鼹鼠),为什么gopher是吉祥物?只能解释为男人的浪漫了。
  3. Go是面向多核的语言,也就是面向未来。Go在Google中也有不少应用,我只找到古老的资料,Rob Pike的演讲Quora的问答、关于dl.google.com故事也很有意思。
  4. Go没有类,但有结构体,我喜欢这点。
  5. Go有指针,太棒了。
  6. Go的interface设计太巧妙了。
  7. Go的goroutine就不用说了,强、优雅。但我目前不知道CSP的细节。
  8. Go的包管理很出色。

体验了2天Go,我无法找到Go的缺点,即使是缺点在我看来也是优点(没有面向对象)。当然Go是否是面向对象的也很难说,可以看看这篇讨论,反正我是认为这是Go的优点。我已经把Go列为第二优先语言了,学了Go可能就不用再学C/C++了。

学习和使用Go的同时能开拓新的思考方式、思维模式,就像学外语一样,思维方式会得到提升。而Go带来的思维改变就是思考要并发。当然并发不是并行,这个可以看Rob Pike的演讲,还有关于并发模式的演讲。形象的说,并发是一个人可以一边吃饭一边喝水,可以随时切换,而并行一般是做同样的事,比如左手和右手同时转笔。我觉得并发有并行的意思,可能之后还要深入理解一下,不过理解了也没啥用。


Go的并发

Go其他部分比较简单,随便学学就行,我觉得并发是Go的精髓。Go的并发很好用很优雅,可以看看性能如何。我将Go与我常用,让我头疼的Python作了对比。首先是裸的并发,就算个加法,也不join,只看起并发的开销。

虽然这个比较不科学,但也可以跑一下试试。

接下来跑个需要同步的,用到了并发常用的queue或channel。

上面两段代码极其相似,这也让我对channel的理解深了一点,功能上就相当于Python中的queue。粗略一测,Go和Python时间性能相差30倍。另外,这段代码中,Python在我的电脑上起不了10000个线程,Go可以起1000000个,我还没测Go的上界。

为什么goroutine这么好用?我搜索了一下,发现很多好的资源,

  1. goroutine和thread的对比,佩服写得好
  2. goroutine不是轻量线程,Rob Pike也说了,理解这点很重要。看了这个说明,完全就理解了

总结

以上是我接触了2天Go写的胡言乱语,希望以后看到不会脸红。总之,Go是不错的语言,值得学习。特别是现在很多情景对并发的需求越来越大、C/C++等语言的设计考虑、硬件的分布式等,Go就是顺应时代潮流的产物,Google还是爸爸。

感叹一句,什么时候我才能那么屌?


参考

  1. Analysis of the Go runtime scheduler,以后看
  2. Five things that make Go fast,很有深度,看了之后脑子炸了,还有一系列

去年与未来

2018年我终于毕业开始工作了。还好当年家里穷,没有选择物理系,否则就与计算机无缘了。一年年的人生经验积累下来,发现自己干的蠢事、令自己回想起来尴尬不已的事、令别人当场窒息的事越来越少。随着年龄的增长,我也发现自己的计算能力慢慢下滑,还好目前的经验增长完全可以抵消算力的下降。2018年我学会了很多技能,一条条来说。


绘画

创造自己理想的二次元对象是我几年来一直想干的事,苦于时间、设备以及其他因素,一直没有付诸行动。我自认为我是绘画天才,尤其对色彩的把控简直天生精准。缺点是,我是写实主义的,追求精确的形与结构,缺乏想象力,创造性不强。为避免老了学得太慢,我在今年年初,终于找到了时间,开始绘画。借助网络、各路大触的教学,我画了一个月,走完了用笔、复印、抄袭、临摹和创作的过程。总共画了10副图,其中有两幅比较有代表性。左图是我情感想画的,右图是我理智想画的。

新子憧 薇尔莉特·伊芙加登紫罗兰

希望老了还有欲望画这样的图,不得不说,我真的是天才。


拍照

由于要去日本玩,所以搞了相机。不过相机是我大三就想买的,苦于没有时间和金钱,一直拖到了今年。我已经按下了7813次快门,其中10%的照片不错到可以给别人看,另20%的照片自己觉得不错,还有10-20%的照片是全景或延迟摄影的额外文件。这样算下来,也就是约50%的快门是浪费的。我不敢说自己是摄影的,只敢说自己是拍照的。因为我的照片都需要经过修改。当然我不会过分修改,P图的原则是改变色彩、保留结构。照片中的人(无意义的人)不属于结构,属于需要抹去的噪点。这么多照片中,有一组我不得不亮出来。

 

这不是说图3有多好看,而是化腐朽为神奇的过程。不要因为表面的难看,而错过内在的好看。希望我能保持这样的眼光,无论对照片、对风景还是对人,要拓展发现美的能力。


日语

我也抽出了不少时间学习日语,虽然差18分考过N2,我已经心满意足,至少我能知道机器翻译错在哪里,目前来说足够了。明年我希望接触韩语,比较简单一点,顺便学习印地语。我深知自己的语言天赋实在没有,跟绘画天赋一比就是负的。不过,不能躺在舒适区,有天赋的要突破极限,没天赋的要达到平均。韩语是2年计划,印地语是20年计划,印度可能就是未来的中国,或许将来需要跟印度人打交道,准备可以先做起来(开个玩笑)。主要是想对比一下,印度文化和中华文明,为什么会产生差距。


视频制作

爱因斯坦会弹小提琴,我语言天赋没有,音乐就不用说了。我只能朝另外的方向发展,别提书法,写的字就像画的。视频是一个方向,视频就是连续的画,也涉及光影、色彩。今年,我做了几个傻视频放抖音,获得了尴尬的点赞,希望明年能熟练运用AE和PR,获得百赞。


 

以上这些技能基本上都跟目前的工作无关,跟工作相关的基本上就是机器翻译。工作在有条不紊地进行,具体也不便说明。团队中的同事知道自己强在哪、弱在哪,这是值得庆幸的。头条还是很不错的厂,我做着自己喜欢的事、自己想做的事,还可以拿钱,感觉人生已经到达了巅峰。

 

毕业本可以拿上海或者北京户口,但是我没有行动,家里有地只是一个原因。如果没有地,我慎重思考之后,基本上也不会放弃我的农村户口。大家都在随波逐流,大家根本没有思考过户籍制度的不合理性。希望在我有生之年,这个等级划分可以取消。目前,我会拿一些不科学的事实安慰自己,江浙人才辈出,北京呢?等。希望我这个随便做的决定,不会在将来引起太大的麻烦。

另外,今年也到了几个地方,日本、海南、杭州、北京、上海、苏州。明年我想去祖国的西边,如成都、重庆、武汉、长沙,日本肯定要再去一次。

毕业之后,身边的好友、打dota的好友换了一波。这是好事也是坏事,证明我不是一滩死水。

希望明年我的博客能月更。