中华网数码

数码
设为书签Ctrl+D将本页面保存为书签,全面了解最新资讯,方便快捷。
业 界/ 互联网/ 行 业/ 通 信/ 数 码/ 手 机/ 平 板/ 笔记本/ 相 机
当前位置: 数码 > 要闻 >

一行代码消除 PyTorch 的 CUDA 内存溢出报错,这个 GitHub 项目刚发布就揽星 600+

一行代码消除 PyTorch 的 CUDA 内存溢出报错,这个 GitHub 项目刚发布就揽星 600+
2021-12-18 22:42:49 来源:量子位

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

CUDA error: out of memory.

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

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

有多厉害?

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

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

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

pip install koila

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

先定义 input、label 和 model:

# A batch of MNIST imageinput = torch.randn(8, 28, 28)# A batch of labelslabel = torch.randn(0, 10, [8])class NeuralNetwork(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),)def forward(self, x):x = self.flatten(x)logits = self.linear_relu_stack(x)return logits

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

loss_fn = CrossEntropyLoss()# Calculate lossesout = nn(t)loss = loss_fn(out, label)# Backward passnn.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 的小哥。

项目地址:点此直达

参考链接:点此直达

责任编辑:kj005

文章投诉热线:156 0057 2229 投诉邮箱:29132 36@qq.com
关键词:

全球5G专利数量排名曝光:华为、LG、三星位列前三

2021-12-16 10:58:23全球5G专利数量排名曝光:华为、LG、三星位列前三

我国仍有1.4亿老年人使用功能机 数字鸿沟问题凸显

2021-12-16 10:39:18我国仍有1.4亿老年人使用功能机 数字鸿沟问题凸显

苹果AirPods系列出货量暴跌 第三季度同比下降33.7%

2021-12-15 10:12:49苹果AirPods系列出货量暴跌 第三季度同比下降33.7%

全球智能手机第三季度出货3.2亿台:三星市场份额第一

2021-12-13 10:44:25全球智能手机第三季度出货3.2亿台:三星市场份额第一

Q3折叠屏手机出货量同比暴涨 环比增长215%

2021-12-10 10:41:47Q3折叠屏手机出货量同比暴涨 环比增长215%

英国2023年将淘汰2G/3G网络 提高5G覆盖率

2021-12-09 10:27:16英国2023年将淘汰2G/3G网络 提高5G覆盖率

相关新闻