1行代碼消除PyTorch的CUDA內(nèi)存溢出報錯，這個GitHub項目攬星600+
本文經(jīng)AI新媒體量子位（公眾號ID:QbitAI）授權(quán)轉(zhuǎn)載，轉(zhuǎn)載請聯(lián)系出處。
多少人用PyTorch“煉丹”時都會被這個bug困擾。
一般情況下，你得找出當(dāng)下占顯存的沒用的程序，然后kill掉。
如果不行，還需手動調(diào)整batch size到合適的大小……
有點麻煩。
現(xiàn)在，有人寫了一個PyTorch wrapper，用一行代碼就能“無痛”消除這個bug。
有多厲害？
相關(guān)項目在GitHub才發(fā)布沒幾天就收獲了600+星。
一行代碼解決內(nèi)存溢出錯誤
軟件包名叫koila，已經(jīng)上傳PyPI，先安裝一下：
- pip install koila
現(xiàn)在，假如你面對這樣一個PyTorch項目：構(gòu)建一個神經(jīng)網(wǎng)絡(luò)來對FashionMNIST數(shù)據(jù)集中的圖像進行分類。
先定義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ù)、計算輸出和losses。
- loss_fn = CrossEntropyLoss()
- # Calculate losses
- out = nn(t)
- loss = loss_fn(out, label)
- # Backward pass
- nn.zero_grad()
- loss.backward()
好了，如何使用koila來防止內(nèi)存溢出？
超級簡單！
只需在第一行代碼，也就是把輸入用lazy張量wrap起來，并指定bacth維度——
koila就能自動幫你計算剩余的GPU內(nèi)存并使用正確的batch size了。
在本例中，batch=0，則修改如下：
- input = lazy(torch.randn(8, 28, 28), batch=0)
完事兒！就這樣和PyTorch“煉丹”時的OOM報錯說拜拜。
靈感來自TensorFlow的靜態(tài)/懶惰評估
下面就來說說koila背后的工作原理。
“CUDA error: out of memory”這個報錯通常發(fā)生在前向傳遞（forward pass）中，因為這時需要保存很多臨時變量。
koila的靈感來自TensorFlow的靜態(tài)/懶惰評估（static/lazy evaluation）。
它通過構(gòu)建圖，并僅在必要時運行訪問所有相關(guān)信息，來確定模型真正需要多少資源。
而只需計算臨時變量的shape就能計算各變量的內(nèi)存使用情況；而知道了在前向傳遞中使用了多少內(nèi)存，koila也就能自動選擇最佳batch size了。
又是算shape又是算內(nèi)存的，koila聽起來就很慢？
NO。
即使是像GPT-3這種具有96層的巨大模型，其計算圖中也只有幾百個節(jié)點。
而Koila的算法是在線性時間內(nèi)運行，任何現(xiàn)代計算機都能夠立即處理這樣的圖計算；再加上大部分計算都是單個張量，所以，koila運行起來一點也不慢。
你又會問了，PyTorch Lightning的batch size搜索功能不是也可以解決這個問題嗎？
是的，它也可以。
但作者表示，該功能已深度集成在自己那一套生態(tài)系統(tǒng)中，你必須得用它的DataLoader，從他們的模型中繼承子類，才能訓(xùn)練自己的模型，太麻煩了。
而koila靈活又輕量，只需一行代碼就能解決問題，非常“大快人心”有沒有。
不過目前，koila還不適用于分布式數(shù)據(jù)的并行訓(xùn)練方法（DDP），未來才會支持多GPU。
以及現(xiàn)在只適用于常見的nn.Module類。
ps. koila作者是一位叫做RenChu Wang的小哥。
項目地址：
https://github.com/rentruewang/koila
