1行代碼消除PyTorch的CUDA內存溢出報錯，這個GitHub項目攬星600+
本文經AI新媒體量子位（公眾號ID:QbitAI）授權轉載，轉載請聯系出處。
多少人用PyTorch“煉丹”時都會被這個bug困擾。
一般情況下，你得找出當下占顯存的沒用的程序，然后kill掉。
如果不行，還需手動調整batch size到合適的大小……
有點麻煩。
現在，有人寫了一個PyTorch wrapper，用一行代碼就能“無痛”消除這個bug。
有多厲害？
相關項目在GitHub才發布沒幾天就收獲了600+星。
一行代碼解決內存溢出錯誤
軟件包名叫koila，已經上傳PyPI，先安裝一下：
- pip install koila
現在，假如你面對這樣一個PyTorch項目：構建一個神經網絡來對FashionMNIST數據集中的圖像進行分類。
先定義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函數、計算輸出和losses。
- loss_fn = CrossEntropyLoss()
- # Calculate losses
- out = nn(t)
- loss = loss_fn(out, label)
- # Backward pass
- nn.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的小哥。
項目地址：
https://github.com/rentruewang/koila
