旷视最近开源了深度学习框架 MegEngine。众所周知,MegDL 一直是贵司的核心技术(大坑)。快速高效,安全可控,灵活定制一直是 MegDL 的核心。老矿工肯定熟悉下面的画面,这是 MegDL 的官方文档主页:

MegDL 官方文档

可以从上面的截图就能非常明显的看出,要简单概括 MegDL,那就是“糙、快、猛”……文档不全,天坑很多,为了业务 xjb 改的地方非常非常多……所以开源其实是一件好事,至少能整理一下代码和文档,减少内部的学习成本和开发成本。

前两天旷视研究院校园关系组让我帮忙在学校论坛里吹吹 MegEngine,还给了我一份稿子:

据说,这个刚刚开源的深度学习框架把LocalConv和sublinear memory做出了新境界!还能让你玩转分布式多机多卡训练,就问你香不香?!
我们不吹不黑地来说一说,事情发生在3月25号,深度学习框架MegEngine开源发布会上,旷视大咖位人物唐文斌、孙剑悉数登场,还有MegEngine负责人田忠博,一起祭出了既可以训得快、训得动、训得爽(一训三得)又可以训推一体、动静合一、灵活高效、兼容并包(四位一体)的工业级框架MegEngine。
说大词你可能体会不到它小而美的细节和优点:
高效实现LocalConv
高效实现LocalConv的前提下,支持任意channel输入输出,
甚至还支持GroupLocalConv(是的,你没有看错)
(PyTorch没有官方支持sublinear memory)
(TensorFlow没有官方支持sublinear memory)
(MXNet 仅半官方支持sublinear memory)
MegEngine官方支持sublinear memory
MegEngine官方支持sublinear memory,还有工业级验证
MegEngine sublinear memory支持任意图
MegEngine sublinear memory是全自动的
.........
那么,欢迎来证实/证伪以上的小细节,也欢迎发现和补充更多的小细节,以下附上传送门:
website:https://megengine.org.cn/
github:https://github.com/MegEngine/MegEngine
知乎:如何看待3月25号开源的旷视深度学习框架天元(MegEngine)?

我看完之后有一些想法:

多机训练是个好东西,但是现在好像还是百花齐放(xjb搞)的阶段吧……

一直在强调的推训一体,我一直没太明白,PyTorch 和 TensorFlow 难道训练和推理要写两套网络么……

动静结合是个好东西,但是下面反复强调的 sublinear memory 和 LocalConv 又是啥玩意呢?

我去搜集了一些资料,我们先从 sublinear memory 讲起:

Training Deep Nets with Sublinear Memory Cost

所谓的 sublinear memory 指的是使用低于 $O(n)$ 的空间复杂度,完成神经网络的训练。

节省内存的关键在于及时回收原地操作。及时回收指的是一些不再需要的中间结果可以丢弃,优先保障当前必要的内存。原地操作指的是输出值覆盖存储在输入值的位置,如果需要输入值,那么进行反操作计算就可以,是时间换空间的思想。

Fig 1. 共享内存。

Fig 1

Fig 2. 及时回收。

Fig 2

算法主要是分块的思想,我们保存每个块的结尾状态而舍弃中间状态。如果需要查询中间状态,那么需要从最近保存的状态进行操作或者反操作,直到中间所需的状态。明显的,这样的空间复杂度就是 $O(n/k+k)$,其中 $k$ 是分块个数。

Alg 1. 从最近状态计算中间状态。

Alg 1

那么遇到下一个问题,如何分块呢?如果我们把很多卷积层分成一块,只保存一个状态,那么想要还原中间某个卷积层的状态是很困难的。所以我们应该把一层卷积与一些 bn,relu,pooling 等分为一块,因为 bn,relu,pooling 的计算开销很小,便于我们还原状态。

具体的,我们在分块时,可以提前给定预期的开销 $B$,然后贪心的最大化每块所需的开销即可。这样可以保证开销低于 $B$ 的情况下,保存了尽可能多的状态。当然了,可以指定块与块直接的分割操作的种类,只在可分割的操作种类上进行分割。我们这里的贪心算法是在遍历计算图,所以分块肯定是在图上的分,需要按照拓扑序遍历,直接递归的搜索或者预处理好拓扑序后 for 一遍都无所谓。

Alg 2. 计算拓扑序。

Alg 2

Alg 3. 贪心的分块。

Alg 3

LocalConv

这是 Tim 老师谈起 LocalConv,老实说我根本没听过这个东西,看样子是很古老的遗留物……

高效的、灵活的 LocalConv/GroupLocalConv

做人脸识别的同学应该知道 LocalConv 这个从 cuda-convnet 年代流传下来的祖传 operator 有多重要。 megengine 的实现在维持高效的情况下,不仅支持任意 channel 输入输出(并没有 channel 必须是 32 的倍数这样的限制),甚至还支持 Group 版的 LocalConv。

开源出来真是业界良心

作者:周昕宇
链接:https://www.zhihu.com/question/377416272/answer/1104186509
来源:知乎

根据我不靠谱的调查发现,LocalConv 是指在 Conv 上分块。以前是一个输入通道对应一个 kernel,现在一个输入通道对应多个 kernel,每个 kernel 都有自己的作用范围。比如说用 4 个 kernel,那么左上角的像素在计算时用一个 kernel,右上角对应一个,左下角对应一个,右下角对应一个。这样每个像素都只用一个 kernel,但是不同块之间的像素用了不同的 kernel……

总结

MegEngine 就是分块大师(暴论

文章目录