一行代码消除PyTorch的CUDA内存溢出报错，这个GitHub项目刚发布就揽星600
多少人用 PyTorch＂炼丹＂时都会被这个 bug 困扰。CUDAerror:outofmemory.
一般情况下，你得找出当下占显存的没用的程序，然后 kill 掉。如果不行，还需手动调整 batch size 到合适的大小，有点麻烦。
现在，有人写了一个 PyTorch wrapper，用一行代码就能＂无痛＂消除这个 bug 。
有多厉害？
相关项目在 GitHub 才发布没几天就收获了 600 + 星。
一行代码解决内存溢出错误
软件包名叫 koila，已经上传 PyPI，先安装一下：pipinstallkoila
现在，假如你面对这样一个 PyTorch 项目：构建一个神经网络来对 FashionMNIST 数据集中的图像进行分类。
先定义 input、label 和 model：#AbatchofMNISTimage input=torch.randn(8,28,28) #Abatchoflabels label=torch.randn(0,10,[8]) 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() #Calculatelosses out=nn(t) loss=loss_fn(out,label) #Backwardpass nn.zero_grad() loss.backward()
好了，如何使用 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 的小哥。
项目地址：点此直达
参考链接：点此直达