本文经AI新媒体量子位（公众号ID：QbitAI）授权转载。转载请联系出处。有多少人在用PyTorch“炼金”的时候会被这个bug困扰。一般来说，你要找出当前占用显存的无用程序，然后杀掉它。如果不行，需要手动调整batchsize到合适的大小……有点麻烦。现在，有人编写了一个PyTorch包装器，用一行代码“无痛”地消除了这个错误。它有多强大？相关项目在GitHub上发布后的几天内获得了600+star。一行代码解决内存溢出错误软件包名为koila，已上传至PyPI。首先安装它：pipinstallkoila现在，假设你面临这样一个PyTorch项目：构建一个神经网络来对FashionMNIST数据集中的图像进行分类。首先定义输入、标签和模型：#AbatchofMNISTimageinput=torch.randn(8,28,28)#Abatchoflabelslabel=torch.randn(0,10,[8])classNeuralNetwork(Module):def__init__(self):super(NeuralNetwork,self).__init__()self.flatten=Flatten()self.linear_relu_stack=Sequential(线性(28*28,512),ReLU(),线性(512,512),ReLU(),线性(512,10),)defforward(self,x):x=self.flatten(x)logits=self.linear_relu_stack(x)returnlogits然后定义损失函数，计算输出和损失。loss_fn=CrossEntropyLoss()#Calculatelossesout=nn(t)loss=loss_fn(out,label)#Backwardpassnn.zero_grad()loss.backward()那么，koila如何防止内存溢出呢？超级简单！你只需要在第一行代码中用惰性张量包裹输入，并指定bacth维度——koila可以自动计算剩余的GPU内存并为你使用正确的批大小。本例中batch=0，那么修改如下：input=lazy(torch.randn(8,28,28),batch=0)大功告成！告别PyTorch“炼金术”时报的OOM错误吧。受TensorFlow的static/lazyevaluation的启发下面说说koila背后的工作原理。“CUDAerror:outofmemory”错误通常发生在forwardpass中，因为此时需要保存很多临时变量。koila的灵感来自于TensorFlow的静态/惰性求值。它通过构建图形并仅在需要访问所有相关信息时运行它来确定模型真正需要多少资源。只需计算临时变量的形状即可计算出每个变量的内存占用情况；并且知道前向传播中使用了多少内存，koila可以自动选择最佳批量大小。它计算形状和记忆，koila听起来很慢？不。即使是像GPT-3这样有96层的巨大模型，其计算图中也只有几百个节点。Koila的算法以线性时间运行，任何现代计算机都可以立即处理此类图计算；再加上大部分计算都是单张量，所以koila一点也不慢。你又会问了，PyTorchLightning的batchsizesearch功能不也能解决这个问题吗？是的，它也可以。但作者表示，这个功能已经深度融入了自己的生态系统。你必须使用它的DataLoader继承他们模型的子类来训练你自己的模型，太麻烦了。而且koila灵活轻量，一行代码就能解决问题，很“幸福”不是吗？不过目前koila还不适用于分布式数据并行训练方式（DDP），未来会支持多GPU。现在只适用于常见的nn.Module类。附言。考拉的作者是一个叫王仁初的小弟弟。项目地址：https://github.com/rentruewang/koila
