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

CUDAerror:outofmemory.

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

现在，有人写了一个PyTorchwrapper，用一行代码就能“无痛”消除这个bug。

有多厉害？

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

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

软件包名叫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(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_grad()loss.backward()

好了，如何使用koila来防止内存溢出？

超级简单！

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

在本例中，batch=0，则修改如下：

input=lazy(torch.randn(8,28,28),batch=0)

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

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

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

“CUDAerror:outofmemory”这个报错通常发生在前向传递（forwardpass）中，因为这时需要保存很多临时变量。

koila的灵感来自TensorFlow的静态/懒惰评估（static/lazyevaluation）。

它通过构建koila作者是一位叫做RenChuWang的小哥。

项目地址：点此直达

参考链接：点此直达

以上就是爱惜日网»一行代码消除PyTorch的CUDA内存溢出报错，这个GitHub项目刚发布就揽星600+的相关内容了，更多精彩请关注作者：爱惜日知识

声明：本文由爱惜日网【创业者资源平台】作者编辑发布，更多技术关注爱惜日技术！