MegEngine 1.4.0 RC1新增API DTR体验贴

在MegEngine的1.4.1 RC1中新增加了DTR的实现,今天抽空试用了一下,所以在此分享一下使用体验。根据官方文档,DTR技术的引入主要是为我们尽可能的节省显存,从而可以使用较大的BatchSize,享受大BatchSize带来的快乐。而DTR的使用方法也比较简单,只需要在训练代码前加入以下两行代码即可,这里需要注意的是memory_budget的设置,目前来说将memory_budget设置为可用显存的一半(例如可用为6GB左右,设置为3GB)就可以正常工作,并且在在使用设定大小的BatchSize会出现OOM的情况下DTR才能发挥出它真正的功效。

from megengine.utils.dtr import DTR
dtr = DTR(memory_budget=int(3 * 1024 ** 3))
#训练代码

DTR论文地址

看来还会有一个 1.4 正式版体验帖 :laughing: