工程Tricks | PyTorch有什么節(jié)省顯存的小技巧?
來源 | 知乎問答
地址?|?https://www.zhihu.com/question/274635237
本文僅作學(xué)術(shù)分享，若侵權(quán)請聯(lián)系后臺刪文處理
回答一：作者-朱小霖
節(jié)省顯存方面，歡迎關(guān)注我們團(tuán)隊最近開源的工作：https://github.com/Tencent/PatrickStar
個人認(rèn)為這個工作把單卡訓(xùn)練，或者是數(shù)據(jù)并行下的顯存節(jié)省做到極致了~
這里主要介紹一下單機(jī)訓(xùn)練上的思路。
隨著模型越來越大，GPU 逐漸從一個計算單元變成一個存儲單元了，顯存的大小限制了能夠訓(xùn)練的模型大小。微軟的 DeepSpeed 團(tuán)隊提出我們其實可以把優(yōu)化器狀態(tài)（Adam 的 momentum 和 variance）放在 CPU 上，用一個實現(xiàn)的比較快的 CPU Adam 來做更新，這樣既不會變慢很多，也可以明顯省出來很多空間。我們把這個思想再往前推一步，我們是不是可以只把需要計算的模型參數(shù)放在 GPU 上，其余的模型參數(shù)，優(yōu)化器狀態(tài)都放在 CPU 上，這樣就可以盡最大能力降低對顯存的需求，讓 GPU 回歸它計算單元的本色。
為了達(dá)成這樣的效果，我們就需要一個動態(tài)的顯存調(diào)度——相對于 DeepSpeed 在訓(xùn)練前就規(guī)定好哪些放在 CPU 上，哪些放在 GPU 上，我們需要在訓(xùn)練過程中實時把下一步需要的模型參數(shù)拿到 GPU 上來。利用 pytorch 的 module hook 可以讓我們在每個 nn.Module 前后調(diào)用回調(diào)函數(shù)，從而動態(tài)把參數(shù)從 CPU 拿到 GPU，或者放回去。
但是，相信大家能夠想象到，如果每次都運(yùn)行到一個 submodule 前，再現(xiàn)把參數(shù)傳上來，肯定就很慢，因為計算得等著 CPU-GPU 的傳輸。為了解決計算效率的問題，我們提出了 chunk-based management。這是什么意思呢？就是我們把參數(shù)按照調(diào)用的順序存儲在了固定大小的 chunk 里（一般是 64M 左右），讓內(nèi)存/顯存的調(diào)度以 chunk 為單位，第一次想把某個 chunk 中的參數(shù)放到 GPU 來的時候，就會直接把整個 chunk 搬到 GPU，這意味著雖然這一次的傳輸可能需要等待，但是在計算下一個 submodule 的時候，因為連著的 module 的參數(shù)都是存在一個 chunk 里的，這些參數(shù)已經(jīng)被傳到 GPU 上來了，從而實現(xiàn)了 prefetch，明顯提升了計算效率。同時，因為 torch 的 allocator 會緩存之前分配的顯存，固定大小的 chunk 可以更高效利用這一機(jī)制，提升顯存利用效率。
在 chunk 的幫助下，我們的模型可以做到，CPU 內(nèi)存加 GPU 顯存有多大，模型就能訓(xùn)多大。和 DeepSpeed 相比，在同等環(huán)境下模型規(guī)?？梢蕴嵘?50%，計算效率（Tflops）也更高。
對于多卡的數(shù)據(jù)并行場景，我們擴(kuò)展了上述方法，可以做到多卡中只有 1 整份模型，1 整份優(yōu)化器狀態(tài)，同時具備了數(shù)據(jù)并行的易用性和模型并行的顯存使用效率。如果想了解多卡訓(xùn)練的方案，以及更詳細(xì)的一些優(yōu)化，也歡迎來看看我們的論文：https://arxiv.org/abs/2108.05818
回答二：作者-鄭哲東
在不修改網(wǎng)絡(luò)結(jié)構(gòu)的情況下, 有如下操作：
1.同意@Jiaming, ?盡可能使用inplace操作， 比如relu 可以使用 inplace=True 。一個簡單的使用方法，如下：
def?inplace_relu(m):
????classname?=?m.__class__.__name__
????if?classname.find('ReLU')?!=?-1:
????????m.inplace=True
model.apply(inplace_relu)
2.進(jìn)一步，比如ResNet 和 DenseNet 可以將 batchnorm 和relu打包成inplace，在bp時再重新計算。使用到了pytorch新的checkpoint特性，有以下兩個代碼。由于需要重新計算bn后的結(jié)果，所以會慢一些。
https://github.com/gpleiss/efficient_densenet_pytorch
https://github.com/mapillary/inplace_abn
3. 每次循環(huán)結(jié)束時 刪除 loss，可以節(jié)約很少顯存，但聊勝于無?？梢娙缦耰ssue
https://discuss.pytorch.org/t/tensor-to-variable-and-memory-freeing-best-practices/6000/2
4. 使用float16精度混合計算。我用過NVIDIA的apex，很好用，可以節(jié)約將近50%的顯存，但是要小心一些不安全的操作如 mean和sum，溢出fp16。
https://github.com/NVIDIA/apex
補(bǔ)充：最近我也嘗試在我CVPR19的GAN模型中加入fp16的訓(xùn)練，可以從15G的顯存需求降到約10G，這樣大多數(shù)1080Ti等較為常見的顯卡就可以訓(xùn)練了。歡迎大家star一波 https://github.com/NVlabs/DG-Net
5. 對于不需要bp的forward，如validation 請使用 torch.no_grad , ?注意model.eval() 不等于 torch.no_grad() 請看如下討論。
https://discuss.pytorch.org/t/model-eval-vs-with-torch-no-grad/19615
6. torch.cuda.empty_cache() 這是del的進(jìn)階版，使用nvidia-smi 會發(fā)現(xiàn)顯存有明顯的變化。但是訓(xùn)練時最大的顯存占用似乎沒變。大家可以試試。
https://discuss.pytorch.org/t/how-can-we-release-gpu-memory-cache/14530
另外，會影響精度的騷操作還有:
把一個batchsize=64分為兩個32的batch，兩次forward以后，backward一次。但會影響 batchnorm等和batchsize相關(guān)的層。
相關(guān)鏈接：老外寫的提高pytorch效率的方法，包含data prefetch等
https://sagivtech.com/2017/09/19/optimizing-pytorch-training-code/
03回答三：作者-Lyken
咦，大家都沒看過陳天奇的 Training Deep Nets with Sublinear Memory Cost 嗎？
訓(xùn)練 CNN 時，Memory 主要的開銷來自于儲存用于計算 backward 的 activation，一般的 workflow 是這樣的
Vanilla backprop
對于一個長度為 N 的 CNN，需要 O(N) 的內(nèi)存。這篇論文給出了一個思路，每隔 sqrt(N) 個 node 存一個 activation，中需要的時候再算，這樣顯存就從 O(N) 降到了 O(sqrt(N))。
Checkpointed backprop
對于越深的模型，這個方法省的顯存就越多，且速度不會明顯變慢。
PyTorch 我實現(xiàn)了一版，有興趣的同學(xué)可以來試試 https://github.com/Lyken17/pytorch-memonger
往期精彩：