多少人用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+的相关内容了 , 更多精彩请关注作者:爱惜日知识
声明:本文由爱惜日网【创业者资源平台】作者编辑发布 , 更多技术关注爱惜日技术!
【一行代码消除PyTorch的CUDA内存溢出报错,这个GitHub项目刚发布就揽星600+】
推荐阅读
- 时空伴随者建康码成了黄码,14天后能消除吗?
- 出现黄码要多少天才会消除
- 女生腋毛浓密怎么消除 女生长腋毛怎么去除
- 行程码怎么消除去过的城市
- 激光脱毛对狐臭有消除吗
- 杨梅酒有酒味怎么消除
- 砭石床垫的作用与功效鼻子里面长痘痘怎么消除
- 下飞机后耳鸣怎样消除
- 行程码带星什么时候能消除
- 《无主之地:前传》全角色传奇模组代码一览