jittor 算子融合方法

jittor简述

由清华大学研制并开源了第一个我国高校自主的深度学习框架——计图(Jittor)。计图是一个完全动态编译(Just-in-time),基于元算子融合和统一计算图的深度学习框架。计图支持30多种的骨干网络,并且开源了多个模型库:对抗生成网络、图像语义分割、检测与实例分割、点云分类、可微渲染等一个完全即时(JIT)编译的深度学习框架。通过JIT编译,我们可以实现更高的性能,同时使系统高度可定制。

2.jittor中的算子融合

Listing 1 Python implementation of convolution using three operators: reindex, broadcast, and sum
1: def conv(x, p):
2:       N,C,H,W = x.shape
3:       o,i,h,w = p.shape
4:       xx = x.reindex(
5:            shape=(N,o,H,W,i,h,w),
6:           indices=("i0", "i4", "i2-i5", "i3-i6")
7:      )
8:       pp = p.broadcast(xx.shape, dims=(0,2,3))
9:       yy = xx*pp
10:      y = yy.sum(dims=(4,5,6))
11:      return y

jittor中的算子融合属于元算子融合,这个例子展示了元算子运算符实现卷积运算。
代码解读:
第1行显示conv有两个参数:x是图像张量,p是参数张量。
第2行和第3行解包了关于参数形状的信息。图像张量x的布局为:批次数(N)、通道数(C)、图像高度(H)和图像宽度(W)。
参数张量p的布局为:输出通道数(o)、输入通道数(I)、内核高度(h)和内核宽度(w)。
第4–7行使用输入张量x和输出张量xx调用reindex运算符。结果是:

xx(i0,i1,i2,i3,i4,i5,i6) = x(i0,i4,i2–i5,i3–i6)

第8行广播参数张量p输出张量pp,形状与xx相同。广播操作符是reindex操作符的专门化,相当于pp = p.reindex(x.shape,indexs =(i1,i4,i5,i6))。这个广播运营商的结果是: pp(i0,i1,i2,i3,i4,i5,i6) = p(i1,i4,i5,i6)

第9行对结果yy(i0,i1,i2,i3,i4,i5,i6) = xx(i0,i1,i2,i3,i4,i5,i6)pp(i0,i1,i2,i3,i4,i5,i6)执行逐元素乘法。
第10行使用sum运算符(reindex-reduce的专门化)来计算y(i0,i1,i2,i3) = X i4,i5,i6 yy(i0,i1,i2,i3,i4,i5,i6)。

上面的例子展示了如何通过对元运算符的4次调用来实现卷积。Jittor能够将这4个元操作符融合成一个操作符,这样中间变量xx、pp、yy就不需要实际计算了。融合所有4个元算子得到最终表达式:y(i0,i1,i2,i3) = X i4,i5,i6 x(i0,i4,I2–i5,i3–i6)p(i1,i4,i5,i6)。以类似的方式,元算子也可以用于实现各种卷积变体,例如膨胀卷积和群卷积。

3. Operator fuser

算子融合是Jittor后端的重要组成部分。它负责任意计算图中的算子融合优化。在上面,我们展示了一个使用卷积计算的算子融合的例子。在实际应用中,前端产生的计算图要复杂得多。为了优化任意情况,我们将计算图视为顶点和边的有向无环图,G = (V,E),其中每个节点V代表一个算子,而每个边E代表一个变量。我们希望将G划分为多个子图Gi'⊆G,其中每个子图Gi' =(Vi',Ei')代表一个融合算子,每个节点恰好属于一个子图,每个边可以属于一个子图或链接两个子图。目标是选择一个执行所有子图的成本最小的分区。然而,准确预测实际执行成本是不可行的:它们取决于硬件和其他因素的许多方面。因此,我们使用一种简化的方法,通过将成本定义为来确定子图

成本定义

其中We, 简单地是由边 e. Eq表示的变量的大小。这个式子对链接两个不同子图的每条边e的权重求和,因此不属于任何子图Gi',这个代价相当于读写指令的总数。这种方法是合理的,因为大多数深度学习模型都受到内存带宽的限制。融合可以通过减少内存操作来提高性能。在最小化成本的同时,需要满足以下规则:
规则1。重新索引操作符不能与前面的元操作符融合,因为这种融合通常会导致性能下降。
规则2。Reindex-reduce运算符不能与以下元运算符融合。这种融合不会提高性能。
规则3。融合不应该在子图之间创建有向循环。例如,给定一个有三个节点和三条边的图:(1 → 2),(2 → 3),(1 → 3),如果第三条边被融合,它将在结果中的子图(1,3)和(平凡的)子图2之间产生一个循环:(1,3) ⇌ 2。

4.最小化搜索成本

使用贪婪算法来最小化成本:在每次迭代中,我们选择满足规则1–3的边e = (vstart,vend),并将vstart,vend融合到vstart所属的子图G’中,重复直到找不到满足规则1–3的边。在实践中,使用动态规划标记算法来避免重复搜索满足规则的边。该算法运行良好,在大多数神经网络中取得了竞争性能。


经典网络组合的融合操作,卷积-归一化-激活。注意操作符是如何在卷积、归一化和激活层之间融合的.png

图显示了经典网络组合的操作融合,卷积-归一化-激活。卷积层由两个重新索引操作符组成,一个逐元素操作符和一个重新索引缩减操作符。规范化层由一个重新索引操作符、一个逐元素操作符和一个重新索引-缩减操作符组成。激活层由多个元素操作符组成。在这种情况下,操作符可以跨卷积、归一化和激活层进行融合。

©著作权归作者所有,转载或内容合作请联系作者
  • 序言:七十年代末,一起剥皮案震惊了整个滨河市,随后出现的几起案子,更是在滨河造成了极大的恐慌,老刑警刘岩,带你破解...
    沈念sama阅读 215,384评论 6 497
  • 序言:滨河连续发生了三起死亡事件,死亡现场离奇诡异,居然都是意外死亡,警方通过查阅死者的电脑和手机,发现死者居然都...
    沈念sama阅读 91,845评论 3 391
  • 文/潘晓璐 我一进店门,熙熙楼的掌柜王于贵愁眉苦脸地迎上来,“玉大人,你说我怎么就摊上这事。” “怎么了?”我有些...
    开封第一讲书人阅读 161,148评论 0 351
  • 文/不坏的土叔 我叫张陵,是天一观的道长。 经常有香客问我,道长,这世上最难降的妖魔是什么? 我笑而不...
    开封第一讲书人阅读 57,640评论 1 290
  • 正文 为了忘掉前任,我火速办了婚礼,结果婚礼上,老公的妹妹穿的比我还像新娘。我一直安慰自己,他们只是感情好,可当我...
    茶点故事阅读 66,731评论 6 388
  • 文/花漫 我一把揭开白布。 她就那样静静地躺着,像睡着了一般。 火红的嫁衣衬着肌肤如雪。 梳的纹丝不乱的头发上,一...
    开封第一讲书人阅读 50,712评论 1 294
  • 那天,我揣着相机与录音,去河边找鬼。 笑死,一个胖子当着我的面吹牛,可吹牛的内容都是我干的。 我是一名探鬼主播,决...
    沈念sama阅读 39,703评论 3 415
  • 文/苍兰香墨 我猛地睁开眼,长吁一口气:“原来是场噩梦啊……” “哼!你这毒妇竟也来了?” 一声冷哼从身侧响起,我...
    开封第一讲书人阅读 38,473评论 0 270
  • 序言:老挝万荣一对情侣失踪,失踪者是张志新(化名)和其女友刘颖,没想到半个月后,有当地人在树林里发现了一具尸体,经...
    沈念sama阅读 44,915评论 1 307
  • 正文 独居荒郊野岭守林人离奇死亡,尸身上长有42处带血的脓包…… 初始之章·张勋 以下内容为张勋视角 年9月15日...
    茶点故事阅读 37,227评论 2 331
  • 正文 我和宋清朗相恋三年,在试婚纱的时候发现自己被绿了。 大学时的朋友给我发了我未婚夫和他白月光在一起吃饭的照片。...
    茶点故事阅读 39,384评论 1 345
  • 序言:一个原本活蹦乱跳的男人离奇死亡,死状恐怖,灵堂内的尸体忽然破棺而出,到底是诈尸还是另有隐情,我是刑警宁泽,带...
    沈念sama阅读 35,063评论 5 340
  • 正文 年R本政府宣布,位于F岛的核电站,受9级特大地震影响,放射性物质发生泄漏。R本人自食恶果不足惜,却给世界环境...
    茶点故事阅读 40,706评论 3 324
  • 文/蒙蒙 一、第九天 我趴在偏房一处隐蔽的房顶上张望。 院中可真热闹,春花似锦、人声如沸。这庄子的主人今日做“春日...
    开封第一讲书人阅读 31,302评论 0 21
  • 文/苍兰香墨 我抬头看了看天上的太阳。三九已至,却和暖如春,着一层夹袄步出监牢的瞬间,已是汗流浃背。 一阵脚步声响...
    开封第一讲书人阅读 32,531评论 1 268
  • 我被黑心中介骗来泰国打工, 没想到刚下飞机就差点儿被人妖公主榨干…… 1. 我叫王不留,地道东北人。 一个月前我还...
    沈念sama阅读 47,321评论 2 368
  • 正文 我出身青楼,却偏偏与公主长得像,于是被迫代替她去往敌国和亲。 传闻我的和亲对象是个残疾皇子,可洞房花烛夜当晚...
    茶点故事阅读 44,248评论 2 352

推荐阅读更多精彩内容