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的边。在实践中,使用动态规划标记算法来避免重复搜索满足规则的边。该算法运行良好,在大多数神经网络中取得了竞争性能。
图显示了经典网络组合的操作融合,卷积-归一化-激活。卷积层由两个重新索引操作符组成,一个逐元素操作符和一个重新索引缩减操作符。规范化层由一个重新索引操作符、一个逐元素操作符和一个重新索引-缩减操作符组成。激活层由多个元素操作符组成。在这种情况下,操作符可以跨卷积、归一化和激活层进行融合。