深度学习 占用gpu内存 使用率为0_深度解析MegEngine亚线性显存优化技术
深度神經網絡訓練是一件復雜的事情，它體現為模型的時間復雜度和空間復雜度，分別對應著計算和內存；而訓練時內存占用問題是漂浮在深度學習社區上空的一塊烏云，如何撥云見日，最大降低神經網絡訓練的內存占用，是一個繞不開的課題。
GPU 顯卡等硬件為深度學習提供了必需的算力，但硬件自身有限的存儲，限制了可訓練模型的尺寸，尤其是大型深度網絡，由此誕生出一系列相關技術，比如亞線性顯存優化、梯度累加、混合精度訓練、分布式訓練，進行 GPU 顯存優化。
其中，亞線性顯存優化方法 [1] 由于較高的計算 / 顯存性價比備受關注；曠視基于此，經過工程擴展和優化，發展出加強版的 MegEngine 亞線性顯存優化技術，輕松把大模型甚至超大模型裝進顯存，也可以毫無壓力使用大 batch 訓練模型。
這里將圍繞著深度學習框架 MegEngine 亞線性顯存優化技術的工程實現和實驗數據，從技術背景、原理、使用、展望等多個方面進行首次深入解讀。
背 ?景在深度學習領域中，隨著訓練數據的增加，需要相應增加模型的尺寸和復雜度，進行模型「擴容」；而 ResNet [2] 等技術的出現在算法層面掃清了訓練深度模型的障礙。不斷增加的數據和持續創新的算法給深度學習框架帶來了新挑戰，能否在模型訓練時有效利用有限的~~ 計算~~ 存儲資源，尤其是減少 GPU 顯存占用，是評估深度學習框架性能的重要指標。
在計算存儲資源一定的情況下，深度學習框架有幾種降低顯存占用的常用方法，其示例如下：
通過合適的梯度定義，讓算子的梯度計算不再依賴于前向計算作為輸入，從而 in-place 地完成算子的前向計算，比如 Sigmoid、Relu 等；? ?
在生命周期沒有重疊的算子之間共享顯存；? ?
通過額外的計算減少顯存占用，比如利用梯度檢查點重新計算中間結果的亞線性顯存優化方法 [1]；? ?
通過額外的數據傳輸減少顯存占用，比如把暫時不用的數據從 GPU 交換到 CPU，需要時再從 CPU 交換回來。
上述顯存優化技術在 MegEngine 中皆有不同程度的實現，這里重點討論基于梯度檢查點的亞線性顯存優化技術。
原 ? 理一個神經網絡模型所占用的顯存空間大體分為兩個方面：?
1)模型本身的參數。
2)模型訓練臨時占用的空間，包括參數的梯度、特征圖等。其中最大占比是 2)中以特征圖形式存在的中間結果，比如，從示例 [1] 可知，根據實現的不同，從 70% 到 90% 以上的顯存用來存儲特征圖。
這里的訓練過程又可分為前向計算，反向計算和優化三個方面，其中前向計算的中間結果最占顯存，還有反向計算的梯度。第 1)方面模型自身的參數內存占用最小。?
MegEngine 加強版亞線性顯存優化技術借鑒了 [1] 的方法，尤其適用于計算存儲資源受限的情況，比如一張英偉達 2080Ti，只有 11G 的顯存；而更貴的 Tesla V100，最大顯存也只有 32G。
? ?圖 1：亞線性顯存優化原理，其中 (b) 保存了 Relu 結果，實際中 Relu 結果可用 in-place 計算
圖 1(a) 給出了卷積神經網絡的基本單元，它由 Conv-BN-Relu 組成。可以看到，反向計算梯度的過程依賴于前向計算獲取的中間結果，一個網絡需要保存的中間結果與其大小成正比，即顯存復雜度為 O(n)。
本質上，亞線性顯存優化方法是以時間換空間，以計算換顯存，如圖 1(b) 所示，它的算法原理如下：
選取神經網絡中 k 個檢查點，從而把網絡分成 k 個 block，需要注意的是，初始輸入也作為一個檢查點；前向計算過程中只保存檢查點處的中間結果；? ?
反向計算梯度的過程中，首先從相應檢查點出發，重新計算單個 block 需要的中間結果，然后計算 block 內部各個 block 的梯度；不同 block 的中間結果計算共享顯存。這種方法有著明顯的優點，即大幅降低了模型的空間復雜度，同時缺點是增加了額外的計算：? ?
顯存占用從 O(n) 變成 O(n/k)+ O(k)，O(n/k) 代表計算單個節點需要的顯存，O(k) 代表 k 個檢查點需要的顯存， 取 k=sqrt(n)，O(n/k)+ O(k)~O(sqrt(n))，可以看到顯存占用從線性變成了亞線性；? ?
因為在反向梯度的計算過程中需要從檢查點恢復中間結果，整體需要額外執行一次前向計算。
在 [1] 的基礎上，MegEngine 結合自身實踐，做了工程擴展和優化，把亞線性顯存優化方法擴展至任意的計算圖，并結合其它常見的顯存優化方法，發展出一套行之有效的加強版亞線性顯存優化技術。
亞線性優化方法采用簡單的網格搜索(grid search)選擇檢查點，MegEngine 在此基礎上增加遺傳算法，采用邊界移動、塊合并、塊分裂等策略，實現更細粒度的優化，進一步降低了顯存占用。
如圖 2 所示，采用型號為 2080Ti 的 GPU 訓練 ResNet50，分別借助基準、亞線性、亞線性 + 遺傳算法三種顯存優化策略，對比了可使用的最大 batch size。僅使用亞線性優化，batch size 從 133 增至 211，是基準的 1.6x；而使用亞線性 + 遺傳算法聯合優化，batch size 進一步增至 262，較基準提升 2x。
圖 2：三種顯存優化方法優化 batch size 的對比：ResNet50
通過選定同一模型、給定 batch size，可以更好地觀察遺傳算法優化顯存占用的情況。如圖 3 所示，隨著迭代次數的增加，遺傳算法逐漸收斂顯存占用，并在第 5 次迭代之后達到一個較穩定的狀態。
圖 3：遺傳算法收斂示意圖
此外，MegEngine 亞線性優化技術通過工程改良，不再局限于簡單的鏈狀結構和同質計算節點, 可用于任意的計算圖，計算節點也可異質，從而拓展了技術的適用場景；并可配合上述顯存優化方法，進一步降低模型的顯存占用。
實 ?驗MegEngine 基于亞線性顯存技術開展了相關實驗，這里固定 batch size=64，在 ResNet18 和 ResNet50 兩個模型上，考察模型訓練時的顯存占用和計算時間。? ?
如圖 4 所示，相較于基準實現，使用 MegEngine 亞線性顯存技術訓練 ResNet18 時，顯存占用降低 32%， 計算時間增加 24%；在較大的 ReNet50 上，顯存占用降低 40%，計算時間增加 25%。同時經過理論分析可知，模型越大，亞線性顯存優化的效果越明顯，額外的計算時間則幾乎不變。圖 4：MegEngine 亞線性優化技術實驗顯存 / 時間對比：ReNet18/ReNet50
在更大模型 Bert 上實驗數據表明，借助 MegEngine 亞線性顯存技術，顯存占用最高降低 75%，而計算時間僅增加 23%，這與理論分析相一致。有興趣的同學可前往 MegEngine ModeHub 試手更多模型實驗：https://megengine.org.cn/model-hub/。
使 ?用MegEngine 官網提供了亞線性顯存優化技術的使用文檔。當你的 GPU 顯存有限，苦于無法訓練較深、較大的神經網絡模型，或者無法使用大 batch 進一步提升深度神經網絡的性能，抑或想要使 batchwise 算子更加穩定，那么，MegEngine 亞線性顯存優化技術正是你需要的解決方案。
上手 MegEngine 亞線性優化技術非常便捷，無需手動設定梯度檢查點，通過幾個簡單的參數，輕松控制遺傳算法的搜索策略。具體使用時，在 MegEngine 靜態圖接口中調用 SublinearMemoryConfig 設置 trace 的參數 sublinear_memory_config，即可打開亞線性顯存優化：
from megengine.jit import trace, SublinearMemoryConfigconfig = SublinearMemoryConfig()
@trace(symbolic=True, sublinear_memory_config=config)
def train_func(data, label, *, net, optimizer):
...
MegEngine 在編譯計算圖和訓練模型時，雖有少量的額外時間開銷，但會顯著緩解顯存不足問題。下面以 ResNet50 為例，說明 MegEngine 可有效突破顯存瓶頸，訓練 batch size 從 100 最高增至 200：
import osfrom multiprocessing import Process
def train_resnet_demo(batch_size, enable_sublinear, genetic_nr_iter=0):
import megengine as mge
import megengine.functional as F
import megengine.hub as hub
import megengine.optimizer as optim
from megengine.jit import trace, SublinearMemoryConfig
import numpy as np
print(
"Run with batch_size={}, enable_sublinear={}, genetic_nr_iter={}".format(
batch_size, enable_sublinear, genetic_nr_iter
)
)
# 使用 GPU 運行這個例子
assert mge.is_cuda_available(), "Please run with GPU"
try:
# 我們從 megengine hub 中加載一個 resnet50 模型。
resnet = hub.load("megengine/models", "resnet50")
optimizer = optim.SGD(resnet.parameters(), lr=0.1,)
config = None
if enable_sublinear:
config = SublinearMemoryConfig(genetic_nr_iter=genetic_nr_iter)
@trace(symbolic=True, sublinear_memory_config=config)
def train_func(data, label, *, net, optimizer):
pred = net(data)
loss = F.cross_entropy_with_softmax(pred, label)
optimizer.backward(loss)
resnet.train()
for i in range(10):
batch_data = np.random.randn(batch_size, 3, 224, 224).astype(np.float32)
batch_label = np.random.randint(1000, size=(batch_size,)).astype(np.int32)
optimizer.zero_grad()
train_func(batch_data, batch_label, net=resnet, optimizer=optimizer)
optimizer.step()
except:
print("Failed")
return
print("Sucess")
# 以下示例結果在 2080Ti GPU 運行得到，顯存容量為 11 GB
# 不使用亞線性內存優化，允許的 batch_size 最大為 100 左右
p = Process(target=train_resnet_demo, args=(100, False))
p.start()
p.join()
# 報錯顯存不足
p = Process(target=train_resnet_demo, args=(200, False))
p.start()
p.join()
# 使用亞線性內存優化，允許的 batch_size 最大為 200 左右
p = Process(target=train_resnet_demo, args=(200, True, 20))
p.start()
p.join(展 ?望
如上所述，MegEngine 的亞線性顯存優化技術通過額外做一次前向計算，即可達到 O(sqrt(n)) 的空間復雜度。如果允許做更多次的前向計算，對整個網絡遞歸地調用亞線性顯存算法，有望在時間復雜度為 O(n log n) 的情況下，達到 O(log n) 的空間復雜度。
更進一步，MegEngine 還將探索亞線性顯存優化技術與數據并行 / 模型并行、混合精度訓練的組合使用問題，以期獲得更佳的集成效果。最后，在 RNN 以及 GNN、Transformer 等其他類型網絡上的使用問題，也是 MegEngine 未來的一個探索方向。
了解更多信息請查詢：
MegEngine GitHub：https://github.com/MegEngine
MegEngine 官網：https://megengine.org.cn
MegEngine ModelHub：https://megengine.org.cn/model-hub/
參考文獻
Chen, T., Xu, B., Zhang, C., & Guestrin, C. (2016). Training deep nets with sublinear memory cost. arXiv preprint arXiv:1604.06174. ?
He, K., Zhang, X., Ren, S., & Sun, J. (2016). Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition (pp. 770-778).
你也「在看」嗎？?
總結
以上是生活随笔為你收集整理的深度学习 占用gpu内存 使用率为0_深度解析MegEngine亚线性显存优化技术的全部內容，希望文章能夠幫你解決所遇到的問題。
- 上一篇： 代理管家app_亲亲小保社保管家app2
- 下一篇： 计算机视觉招聘_INDEMIND|SLA