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

Python再入门

会用一门语言写代码并不能说是掌握了这门语言,尤其是入门容易的语言,掌握起来更困难。你说你会用汇编写程序,我可以承认你掌握了汇编语言,而你说会用Python写程序,十有八九真的只是会写而已。

有基础的哥们,Python能在3分钟内入门,其他解释性语言也是如此,不用花太多时间,就能说我会了,只要能让我查google。我从大一(2011)开始接触Python,几乎每个月都在学新的内容。当然,我新学的内容不是Pythonic。


Python基础库

我很久以前买了一本很厚的《Python标准库》,结果才看了不到100页。最近,我才真正认识到标准库的作用。此处,我不会考虑Python 2,我只会Python 3。

unicodedata

unicodedata对我这样的多语言工作者还是比较有用的,目前这个库只提供了各个unicode字符所属的类别,如’我’是符号类别(‘Lo’),’1’是数字类别(‘Nd’),这样我就不用翻wikipedia了。不过,这个库也没有完全满足我的需求,经常需要查阅unicode各个块(unicode_block)的信息。

collections

提供了除基础容器list,tuple,dict,set以外的容器,我比较常用的是namedtuple和defaultdict。defaultdict就是在赋值的时候不用手动初始化,考虑简单的计数问题(仅仅是实例,可以用counter)。

复杂一点的,如针对各个字符串统计各字符数的defaultdict,也就是defaultdict的defaultdict。

array

与list类似,不过array的数据存储更加原生,而且array中的数据类型得一致,就像C/C++一样。看例子就明白了。

此处我反正是搞不懂为啥a只用87624个bytes,这乘法不对啊。40064就容易理解多利,10000个整数,每个4 bytes,加上数据结构64 bytes,这个空间直接省了1半。当然也可以用numpy。

linecache

随机读取文件行,小文件可以,大文件直接内存爆炸。

timeit

规范的测试几行代码速度的工具,使用也很方便。

operator

我曾经一直不理解operator的用处,跟内置的运算符重复,实在没有,也可以使用lambda。后来我搜索了一下,stackoverflow,stackoverflow2,给了我一点启示。主要出于速度、代码可读性、以及picklable考虑。关于速度可以简单测试如下,

itertools

与operator一样,属于函数式编程模块。借鉴了Haskell等语言的思想,实现了大量高效、内存友好的迭代工具,与operator相辅相成。官方介绍中也提供了很多例子。我用得最多的居然是chain.from_iterable(),我觉得我得改。takewhile这些都很不错。

functools

与上面两个一样,属于函数式编程模块。其中的lru_cache是比较典型的操作,能很好的解释functools。另外wraps也是常用的,其他几个我还没有使用过。

typing

使用type hints可以帮助我们在编码时就发现bug,当然借助的是ide的提示。mypy也可以进行代码检查,可以告知哪里类型有问题,可以及时发现处理。

其他

其他很多库,如heapq,bisect,re,threading,multiprocessing,profile,pdb,json等都是非常常用好用的库。如果能灵活使用这些库,必定能提供编码效率。此处效率包括编程速度和运行速度。估计我也就掌握了1/10左右。


每个半年回过头看自己写的老代码,总感觉非常羞耻,想推倒重来。我觉得这是好事,说明人还在进步学习。希望还能保持这个速度。

链接

  1. Python 3基础库

 

为何使用Python

        我决定重新好好学学python,重新审视一下自己,我没有一门精通的语言。想javascript,matlab,c/c++,php,shell什么的看似能报出一大堆,我都不敢说精通。我能做的就是用他们来做编程题,而不是做项目。我不了解各种语言的精髓,这样不太好。所以决定从python开始,好好学习。看的是learning python,他还有姐妹篇programing python。选python的原因是python对于安全和逆向工程来说是必备工具。当然还有其他原因。正如作者所说。python能做任何事,从系统管理到图片处理,从嵌入式到机器人,几乎无所不能。python很漂亮,写出的代码就像文档一样。python稳定而强大。

        正如python的哲学:

链接:

  1. python homepage
  2. List Application
  3. Mayavi