1行代碼消除PyTorch的CUDA內(nèi)存溢出報(bào)錯(cuò)，這個(gè)GitHub項(xiàng)目剛發(fā)布就攬星600+

豐色 2021-12-18 17:43:49 來(lái)源：量子位

豐色 發(fā)自 凹非寺
量子位 報(bào)道 | 公眾號(hào) QbitAI

CUDA error: out of memory.

多少人用PyTorch“煉丹”時(shí)都會(huì)被這個(gè)bug困擾。

一般情況下，你得找出當(dāng)下占顯存的沒(méi)用的程序，然后kill掉。

如果不行，還需手動(dòng)調(diào)整batch size到合適的大小……

有點(diǎn)麻煩。

現(xiàn)在，有人寫(xiě)了一個(gè)PyTorch wrapper，用一行代碼就能“無(wú)痛”消除這個(gè)bug。

有多厲害？

相關(guān)項(xiàng)目在GitHub才發(fā)布沒(méi)幾天就收獲了600+星。

一行代碼解決內(nèi)存溢出錯(cuò)誤

軟件包名叫koila，已經(jīng)上傳PyPI，先安裝一下：

pip install koila

現(xiàn)在，假如你面對(duì)這樣一個(gè)PyTorch項(xiàng)目：構(gòu)建一個(gè)神經(jīng)網(wǎng)絡(luò)來(lái)對(duì)FashionMNIST數(shù)據(jù)集中的圖像進(jìn)行分類。

先定義input、label和model：

# A batch of MNIST image
input = torch.randn(8, 28, 28)

# A batch of labels
label = 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函數(shù)、計(jì)算輸出和losses。

loss_fn = CrossEntropyLoss()

# Calculate losses
out = nn(t)
loss = loss_fn(out, label)

# Backward pass
nn.zero_grad()
loss.backward()

好了，如何使用koila來(lái)防止內(nèi)存溢出？

超級(jí)簡(jiǎn)單！

只需在第一行代碼，也就是把輸入用lazy張量wrap起來(lái)，并指定bacth維度——

koila就能自動(dòng)幫你計(jì)算剩余的GPU內(nèi)存并使用正確的batch size了。

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

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

完事兒！就這樣和PyTorch“煉丹”時(shí)的OOM報(bào)錯(cuò)說(shuō)拜拜。

靈感來(lái)自TensorFlow的靜態(tài)/懶惰評(píng)估

下面就來(lái)說(shuō)說(shuō)koila背后的工作原理。

“CUDA error: out of memory”這個(gè)報(bào)錯(cuò)通常發(fā)生在前向傳遞（forward pass）中，因?yàn)檫@時(shí)需要保存很多臨時(shí)變量。

koila的靈感來(lái)自TensorFlow的靜態(tài)/懶惰評(píng)估（static/lazy evaluation）。

它通過(guò)構(gòu)建圖，并僅在必要時(shí)運(yùn)行訪問(wèn)所有相關(guān)信息，來(lái)確定模型真正需要多少資源。

而只需計(jì)算臨時(shí)變量的shape就能計(jì)算各變量的內(nèi)存使用情況；而知道了在前向傳遞中使用了多少內(nèi)存，koila也就能自動(dòng)選擇最佳batch size了。

又是算shape又是算內(nèi)存的，koila聽(tīng)起來(lái)就很慢？

NO。

即使是像GPT-3這種具有96層的巨大模型，其計(jì)算圖中也只有幾百個(gè)節(jié)點(diǎn)。

而Koila的算法是在線性時(shí)間內(nèi)運(yùn)行，任何現(xiàn)代計(jì)算機(jī)都能夠立即處理這樣的圖計(jì)算；再加上大部分計(jì)算都是單個(gè)張量，所以，koila運(yùn)行起來(lái)一點(diǎn)也不慢。

你又會(huì)問(wèn)了，PyTorch Lightning的batch size搜索功能不是也可以解決這個(gè)問(wèn)題嗎？

是的，它也可以。

但作者表示，該功能已深度集成在自己那一套生態(tài)系統(tǒng)中，你必須得用它的DataLoader，從他們的模型中繼承子類，才能訓(xùn)練自己的模型，太麻煩了。

而koila靈活又輕量，只需一行代碼就能解決問(wèn)題，非?！按罂烊诵摹庇袥](méi)有。

不過(guò)目前，koila還不適用于分布式數(shù)據(jù)的并行訓(xùn)練方法（DDP），未來(lái)才會(huì)支持多GPU。

以及現(xiàn)在只適用于常見(jiàn)的nn.Module類。

ps. koila作者是一位叫做RenChu Wang的小哥。

項(xiàng)目地址：
https://github.com/rentruewang/koila

參考鏈接：
https://www.reddit.com/r/MachineLearning/comments/r4zaut/p_eliminate_pytorchs_cuda_error_out_of_memory/