您的位置:首页 > 财经 >

一行代码消除PyTorch的CUDA内存溢出报错,这个GitHub项目刚

来源:IT之家 2021-12-20 02:00   阅读量:6970   

分享

多少人用 PyTorch炼丹时都会被这个 bug 困扰。

一行代码消除PyTorch的CUDA内存溢出报错,这个GitHub项目刚

CUDAerror:outofmemory.

一般情况下,你得找出当下占显存的没用的程序,然后 kill 掉如果不行,还需手动调整 batch size 到合适的大小,有点麻烦

现在,有人写了一个 PyTorch wrapper,用一行代码就能无痛消除这个 bug。

有多厉害。

相关项目在 GitHub 才发布没几天就收获了 600 + 星。

一行代码解决内存溢出错误

软件包名叫 koila,已经上传 PyPI,先安装一下:

pipinstallkoila

现在,假如你面对这样一个 PyTorch 项目:构建一个神经网络来对 FashionMNIST 数据集中的图像进行分类。

先定义 input,label 和 model:

#AbatchofMNISTimageinput=torch.randn#Abatchoflabelslabel=torch.randn)classNeuralNetwork(Module):def__init__(self):super(NeuralNetwork,self).__init__()self.flatten=Flatten()self.linear_relu_stack=Sequential(Linear(28*28,512),ReLU(),Linear(512,512),ReLU(),Linear(512,10),)defforward(self,x):x=self.flatten(x)logits=self.linear_relu_stack(x)returnlogits

然后定义 loss 函数,计算输出和 losses。

loss_fn=CrossEntropyLoss#Calculatelossesout=nn(t)loss=loss_fn(out,label)#Backwardpassnn.zero_gradloss.backward

好了,如何使用 koila 来防止内存溢出。

超级简单!

只需在第一行代码,也就是把输入用 lazy 张量 wrap 起来,并指定 bacth 维度,koila 就能自动帮你计算剩余的 GPU 内存并使用正确的 batch size 了。

在本例中,batch=0,则修改如下:

input=lazy,batch=0)

完事儿!就这样和 PyTorch炼丹时的 OOM 报错说拜拜。

灵感来自 TensorFlow 的静态 / 懒惰评估

下面就来说说 koila 背后的工作原理。

CUDA error: out of memory这个报错通常发生在前向传递中,因为这时需要保存很多临时变量。

koila 的灵感来自 TensorFlow 的静态 / 懒惰评估。

它通过构建图,并仅在必要时运行访问所有相关信息,来确定模型真正需要多少资源。。

而只需计算临时变量的 shape 就能计算各变量的内存使用情况,而知道了在前向传递中使用了多少内存,koila 也就能自动选择最佳 batch size 了。

又是算 shape 又是算内存的,koila 听起来就很慢。

NO。

即使是像 GPT—3 这种具有 96 层的巨大模型,其计算图中也只有几百个节点。

而 Koila 的算法是在线性时间内运行,任何现代计算机都能够立即处理这样的图计算,再加上大部分计算都是单个张量,所以,koila 运行起来一点也不慢。

你又会问了,PyTorch Lightning 的 batch size 搜索功能不是也可以解决这个问题吗。

是的,它也可以。

而 koila 灵活又轻量,只需一行代码就能解决问题,非常大快人心有没有。

不过目前,koila 还不适用于分布式数据的并行训练方法,未来才会支持多 GPU。通过在PyTorchLightning中设置混合精度标志,它将在可能的情况下自动使用半精度,同时在其他地方保留单精度。

以及现在只适用于常见的 nn.Module 类。

项目地址:点此直达

参考链接:点此直达

郑重声明:此文内容为本网站转载企业宣传资讯,目的在于传播更多信息,与本站立场无关。仅供读者参考,并请自行核实相关内容。