报道 | 量子位
多少人用PyTorch“炼丹”时都会被这个bug困扰。 一般情况下,你得找出当下占显存的没用的程序,然后kill掉。 如果不行,还需手动调整batch size到合适的大小…… 有点麻烦。 现在,有人写了一个PyTorch wrapper,用一行代码就能“无痛”消除这个bug。 有多厉害? 相关项目在GitHub才发布没几天就收获了600+星。 一行代码解决内存溢出错误软件包名叫koila,已经上传PyPI,先安装一下:
现在,假如你面对这样一个PyTorch项目:构建一个神经网络来对FashionMNIST数据集中的图像进行分类。 先定义input、label和model: # A batch of MNIST image 然后定义loss函数、计算输出和losses。
好了,如何使用koila来防止内存溢出? 超级简单! 只需在第一行代码,也就是把输入用lazy张量wrap起来,并指定bacth维度—— koila就能自动帮你计算剩余的GPU内存并使用正确的batch size了。 在本例中,batch=0,则修改如下: input = lazy(torch.randn(8, 28, 28), batch=0) 完事儿!就这样和PyTorch“炼丹”时的OOM报错说拜拜。 灵感来自TensorFlow的静态/懒惰评估下面就来说说koila背后的工作原理。 “CUDA error: out of memory”这个报错通常发生在前向传递(forward pass)中,因为这时需要保存很多临时变量。 koila的灵感来自TensorFlow的静态/懒惰评估(static/lazy evaluation)。 它通过构建图,并仅在必要时运行访问所有相关信息,来确定模型真正需要多少资源。 而只需计算临时变量的shape就能计算各变量的内存使用情况;而知道了在前向传递中使用了多少内存,koila也就能自动选择最佳batch size了。 又是算shape又是算内存的,koila听起来就很慢? NO。 即使是像GPT-3这种具有96层的巨大模型,其计算图中也只有几百个节点。 而Koila的算法是在线性时间内运行,任何现代计算机都能够立即处理这样的图计算;再加上大部分计算都是单个张量,所以,koila运行起来一点也不慢。 你又会问了,PyTorch Lightning的batch size搜索功能不是也可以解决这个问题吗? 是的,它也可以。 但作者表示,该功能已深度集成在自己那一套生态系统中,你必须得用它的DataLoader,从他们的模型中继承子类,才能训练自己的模型,太麻烦了。 而koila灵活又轻量,只需一行代码就能解决问题,非常“大快人心”有没有。 不过目前,koila还不适用于分布式数据的并行训练方法(DDP),未来才会支持多GPU。 以及现在只适用于常见的nn.Module类。 ps. koila作者是一位叫做RenChu Wang的小哥。 项目地址: 参考链接: |
|