工程之道：旷视天元框架亚线性显存优化技术解析

2020-05-08 01:29:51

机械之心发布

[好文分享：www.ii77.com]

[转载出处：www.ii77.com]

旷视研究院

基于梯度搜检点的亚线性显存优化方式 [1] 因为较高的较量/显存性价比受到存眷。MegEngine 经由工程扩展和优化，成长出一套行之有效的增强版亚线性显存优化手艺，既可在较量存储资源受限的前提下，轻松练习更深的模型，又可使用更大 batch size，进一步提拔模型机能，不乱 batchwise 算子。使用 MegEngine 练习 ResNet18/ResNet50，显存占用离别最高降低 23%/40%；在更大的 Bert 模型上，降幅更是高达 75%，而额外的较量开销几乎不变。

深度神经收集练习是一件复杂的事情，它施展为模型的时间复杂度和空间复杂度，离别对应着较量和内存；而练习时内存占用问题是漂浮在深度进修社区上空的一块乌云，若何拨云见日，最大降低神经收集练习的内存占用，是一个绕不开的课题。

GPU 显卡等硬件为深度进修供应了必需的算力，但硬件自身有限的存储，限制了可练习模型的尺寸，尤其是大型深度收集，由此降生出一系列相关手艺，好比亚线性显存优化、梯度累加、夹杂精度练习、分布式练习，进行 GPU 显存优化。

个中，亚线性显存优化方式 [1] 因为较高的较量/显存性价比备受存眷；旷视基于此，经由工程扩展和优化，成长出增强版的 MegEngine 亚线性显存优化手艺，轻松把大模型甚至超大模型装进显存，也能够毫无压力使用大 batch 练习模型。

这里将环绕着深度进修框架 MegEngine 亚线性显存优化手艺的工程实现和实验数据，从手艺配景、道理、使用、瞻望等多个方面进行首次深入解读。

配景

在深度进修范畴中，跟着练习数据的增加，需要响应增加模型的尺寸和复杂度，进行模型「扩容」；而 ResNet [2] 等手艺的显现在算法层面扫清了练习深度模型的障碍。络续增加的数据和持续立异的算法给深度进修框架带来了新挑战，可否在模型练习时有效行使有限的较量存储资源，尤其是削减 GPU 显存占用，是评估深度进修框架机能的主要指标。

在较量存储资源必然的情形下，深度进修框架有几种降低显存占用的常用方式，其示例如下：

经由合适的梯度界说，让算子的梯度较量不再依靠于前向较量作为输入，从而 in-place 地完成算子的前向较量，好比 Sigmoid、Relu 等；
在生命周期没有重叠的算子之间共享显存；
经由额外的较量削减显存占用，好比行使梯度搜检点从新较量中央究竟的亚线性显存优化方式 [1]；
经由额外的数据传输削减显存占用，好比把临时不消的数据从 GPU 交流到 CPU，需要时再从 CPU 交流回来。

上述显存优化手艺在 MegEngine 中皆有分歧水平的实现，这里重点商议基于梯度搜检点的亚线性显存优化手艺。

道理

一个神经收集模型所占用的显存空间大体分为两个方面：1）模型自己的参数，2）模型练习暂时占用的空间，包罗参数的梯度、特征图等。个中最大占比是 2）中以特征图形式存在的中央究竟，好比，从示例 [1] 可知，凭据实现的分歧，从 70% 到 90% 以上的显存用来存储特征图。

这里的练习过程又可分为前向较量，反向较量和优化三个方面，个中前向较量的中央究竟最占显存，还有反向较量的梯度。第 1）方面模型自身的参数内存占用最小。

MegEngine 增强版亚线性显存优化手艺借鉴了 [1] 的方式，尤其适用于较量存储资源受限的情形，好比一张英伟达 2080Ti，只有 11G 的显存；而更贵的 Tesla V100，最大显存也只有 32G。

图 1：亚线性显存优化道理，个中 (b) 留存了 Relu 究竟，实际中 Relu 究竟可用 in-place 较量

图 1(a) 给出了卷积神经收集的根基单元，它由 Conv-BN-Relu 构成。能够看到，反向较量梯度的过程依靠于前向较量获取的中央究竟，一个收集需要留存的中央究竟与其巨细成正比，即显存复杂度为 O(n)。

素质上，亚线性显存优化方式是以时间换空间，以较量换显存，如图 1(b) 所示，它的算法道理如下：

拔取神经收集中 k 个搜检点，从而把收集分成 k 个 block，需要注重的是，初始输入也作为一个搜检点；前向较量过程中只留存搜检点处的中央究竟；
反向较量梯度的过程中，首先从响应搜检点出发，从新较量单个 block 需要的中央究竟，然后较量 block 内部各个 block 的梯度；分歧 block 的中央究竟较量共享显存。

这种方式有着显着的长处，即大幅降低了模型的空间复杂度，同时瑕玷是增加了额外的较量：

显存占用从 O(n) 酿成 O(n/k)+ O(k)，O(n/k) 代表较量单个节点需要的显存，O(k) 代表 k 个搜检点需要的显存，取 k=sqrt(n)，O(n/k)+ O(k)~O(sqrt(n))，能够看到显存占用从线性酿成了亚线性；
因为在反向梯度的较量过程中需要从搜检点恢复中央究竟，整体需要额外执行一次前向较量。

工程

在 [1] 的根蒂上，MegEngine 连系自身实践，做了工程扩展和优化，把亚线性显存优化方式扩展至随意的较量图，并连系另外常见的显存优化方式，成长出一套行之有效的增强版亚线性显存优化手艺。

亚线性优化方式采用简洁的网格搜刮（grid search）选择搜检点，MegEngine 在此根蒂上增加遗传算法，采用界限移动、块归并、块盘据等策略，实现更细粒度的优化，进一步降低了显存占用。

如图 2 所示，采用型号为 2080Ti 的 GPU 练习 ResNet50，离别借助基准、亚线性、亚线性+遗传算法三种显存优化策略，对比了可使用的最大 batch size。仅使用亚线性优化，batch size 从 133 增至 211，是基准的 1.6x；而使用亚线性+遗传算法结合优化，batch size 进一步增至 262，较基准提拔 2x。

图 2：三种显存优化方式优化 batch size 的对比：ResNet50

经由选定统一模型、给定 batch size，能够更好地视察遗传算法优化显存占用的情形。如图 3 所示，跟着迭代次数的增加，遗传算法逐渐收敛显存占用，并在第 5 次迭代之后达到一个较不乱的状况。

图 3：遗传算法收敛示意图

此外，MegEngine 亚线性优化手艺经由工程改良，不再局限于简洁的链状构造和同质较量节点, 可用于随意的较量图，较量节点也可异质，从而拓展了手艺的适用场景；并可合营上述显存优化方式，进一步降低模型的显存占用。

实验

MegEngine 基于亚线性显存手艺开展了相关实验，这里固定 batch size=64，在 ResNet18 和 ResNet50 两个模型上，考查模型练习时的显存占用和较量时间。

如图 4 所示，相较于基准实现，使用 MegEngine 亚线性显存手艺练习 ResNet18 时，显存占用降低 32%，较量时间增加 24%；在较大的 ReNet50 上，显存占用降低 40%，较量时间增加 25%。同时经由理论剖析可知，模型越大，亚线性显存优化的结果越显着，额外的较量时间则几乎不变。

图 4：MegEngine 亚线性优化手艺实验显存/时间对比：ReNet18/ReNet50

在更大模型 Bert 上实验数据表明，借助 MegEngine 亚线性显存手艺，显存占用最高降低 75%，而较量时间仅增加 23%，这与理论剖析相一致。有乐趣的同窗可前去 MegEngine ModeHub 试手更多模型实验：https://megengine.org.cn/model-hub/。

使用

MegEngine 官网供应了亚线性显存优化手艺的使用文档。当你的 GPU 显存有限，吃力于无法练习较深、较大的神经收集模型，或许无法使用大 batch 进一步提拔深度神经收集的机能，抑或想要使 batchwise 算子加倍不乱，那么，MegEngine 亚线性显存优化手艺恰是你需要的解决方案。

上手 MegEngine 亚线性优化手艺非常便捷，无需手动设定梯度搜检点，经由几个简洁的参数，轻松掌握遗传算法的搜刮策略。具体使用时，在 MegEngine 静态图接口中挪用 SublinearMemoryConfig 设置 trace 的参数 sublinear_memory_config，即可打开亚线性显存优化：

from megengine.jit import trace, SublinearMemoryConfig
config = SublinearMemoryConfig()
@trace(symbolic=True, sublinear_memory_config=config)def train_func(data, label, *, net, optimizer): ...

MegEngine 在编译较量图和练习模型时，虽有少量的额外时间开销，但会显著缓解显存不足问题。下面以 ResNet50 为例，解说 MegEngine 可有效冲破显存瓶颈，练习 batch size 从 100 最高增至 200：

import osfrom multiprocessing import Process

def train_resnet_demo(batch_size, enable_sublinear, genetic_nr_iter=0):    import megengine as mge    import megengine.functional as F    import megengine.hub as hub    import megengine.optimizer as optim    from megengine.jit import trace, SublinearMemoryConfig    import numpy as np
    print(        "Run with batch_size={}, enable_sublinear={}, genetic_nr_iter={}".format(            batch_size, enable_sublinear, genetic_nr_iter        )    )    # 使用GPU运行这个例子    assert mge.is_cuda_available(), "Please run with GPU"    try:        # 我们从 megengine hub 中加载一个 resnet50 模型。        resnet = hub.load("megengine/models", "resnet50")
        optimizer = optim.SGD(resnet.parameters(), lr=0.1,)
        config = None        if enable_sublinear:            config = SublinearMemoryConfig(genetic_nr_iter=genetic_nr_iter)
        @trace(symbolic=True, sublinear_memory_config=config)        def train_func(data, label, *, net, optimizer):            pred = net(data)            loss = F.cross_entropy_with_softmax(pred, label)            optimizer.backward(loss)
        resnet.train()        for i in range(10):            batch_data = np.random.randn(batch_size, 3, 224, 224).astype(np.float32)            batch_label = np.random.randint(1000, size=(batch_size,)).astype(np.int32)            optimizer.zero_grad()            train_func(batch_data, batch_label, net=resnet, optimizer=optimizer)            optimizer.step()    except:        print("Failed")        return
    print("Sucess")

# 以下示例究竟在2080Ti GPU运行获得，显存容量为 11 GB
# 不使用亚线性内存优化，许可的batch_size最大为 100 摆布p = Process(target=train_resnet_demo, args=(100, False))p.start()p.join()# 报错显存不足p = Process(target=train_resnet_demo, args=(200, False))p.start()p.join()
# 使用亚线性内存优化，许可的batch_size最大为 200 摆布p = Process(target=train_resnet_demo, args=(200, True, 20))p.start()p.join()

瞻望

如上所述，MegEngine 的亚线性显存优化手艺经由额外做一次前向较量，即可达到 O(sqrt(n)) 的空间复杂度。若是许可做更多次的前向较量，对整个收集递归地挪用亚线性显存算法，有望在时间复杂度为 O(n log n) 的情形下，达到 O(log n) 的空间复杂度。

更进一步，MegEngine 还将索求亚线性显存优化手艺与数据并行/模型并行、夹杂精度练习的组合使用问题，以期获得更佳的集成绩果。最后，在 RNN 以及 GNN、Transformer 等其他类型收集上的使用问题，也是 MegEngine 将来的一个索求偏向。

迎接接见

MegEngine GitHub：https://github.com/MegEngine
MegEngine 官网：https://megengine.org.cn
MegEngine ModelHub：https://megengine.org.cn/model-hub/

参考文献

1.Chen, T., Xu, B., Zhang, C., & Guestrin, C. (2016). Training deep nets with sublinear memory cost. arXiv preprint arXiv:1604.06174.

2.He, K., Zhang, X., Ren, S., & Sun, J. (2016). Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition (pp. 770-778).

本文为机械之心发布，转载请关联本公家号获得授权。

✄------------------------------------------------

到场机械之心（全职记者 / 实习生）：hr@jiqizhixin.com

投稿或追求报道：content@jiqizhixin.com

告白 & 商务合作：bd@jiqizhixin.com

工程之道：旷视天元框架亚线性显存优化技术解析

热门文章

小编推荐

推荐文章