[Pytorch]深度模型的視訊記憶體計算以及優化
原文連結：https://oldpan.me/archives/how-to-calculate-gpu-memory
前言
親，視訊記憶體炸了，你的顯示卡快冒煙了！
torch.FatalError: cuda runtime error (2) : out of memory at /opt/conda/conda-bld/pytorch_1524590031827/work/aten/src/THC/generic/THCStorage.cu:58
想必這是所有煉丹師們最不想看到的錯誤，沒有之一。
OUT OF MEMORY，顯然是視訊記憶體裝不下你那麼多的模型權重還有中間變數，然後程式奔潰了。怎麼辦，其實辦法有很多，及時清空中間變數，優化程式碼，減少batch，等等等等，都能夠減少視訊記憶體溢位的風險。
但是這篇要說的是上面這一切優化操作的基礎，如何去計算我們所使用的視訊記憶體。學會如何計算出來我們設計的模型以及中間變數所佔視訊記憶體的大小，想必知道了這一點，我們對自己視訊記憶體也就會得心應手了。
如何計算
首先我們應該瞭解一下基本的資料量資訊：
- 1 G = 1000 MB
- 1 M = 1000 KB
- 1 K = 1000 Byte
- 1 B = 8 bit
好，肯定有人會問為什麼是1000而不是1024，這裡不過多討論，只能說兩種說法都是正確的，只是應用場景略有不同。這裡統一按照上面的標準進行計算。
然後我們說一下我們平常使用的向量所佔的空間大小，以Pytorch官方的資料格式為例（所有的深度學習框架資料格式都遵循同一個標準）：
我們只需要看左邊的資訊，在平常的訓練中，我們經常使用的一般是這兩種型別：
- float32 單精度浮點型
- int32 整型
一般一個8-bit的整型變數所佔的空間為1B也就是8bit。而32位的float則佔4B也就是32bit。而雙精度浮點型double和長整型long在平常的訓練中我們一般不會使用。
ps：消費級顯示卡對單精度計算有優化，伺服器級別顯示卡對雙精度計算有優化。
也就是說，假設有一幅RGB三通道真彩色圖片，長寬分別為500 x 500，資料型別為單精度浮點型，那麼這張圖所佔的視訊記憶體的大小為：500 x 500 x 3 x 4B = 3M。
而一個(256,3,100,100)-(N,C,H,W)的FloatTensor所佔的空間為256 x 3 x 100 x 100 x 4B = 31M
不多是吧，沒關係，好戲才剛剛開始。
視訊記憶體去哪兒了
看起來一張圖片(3x256x256)和卷積層(256x100x100)所佔的空間並不大，那為什麼我們的視訊記憶體依舊還是用的比較多，原因很簡單，佔用視訊記憶體比較多空間的並不是我們輸入影象，而是神經網路中的中間變數以及使用optimizer演算法時產生的巨量的中間引數。
我們首先來簡單計算一下Vgg16這個net需要佔用的視訊記憶體：
通常一個模型佔用的視訊記憶體也就是兩部分：
- 模型自身的引數(params)
- 模型計算產生的中間變數(memory)
圖片來自cs231n，這是一個典型的sequential-net，自上而下很順暢，我們可以看到我們輸入的是一張224x224x3的三通道影象，可以看到一張影象只佔用150x4k，但上面標註的是150k，這是因為上圖中在計算的時候預設的資料格式是8-bit而不是32-bit，所以最後的結果要乘上一個4。
我們可以看到，左邊的memory值代表：影象輸入進去，圖片以及所產生的中間卷積層所佔的空間。我們都知道，這些形形色色的深層卷積層也就是深度神經網路進行“思考”的過程：
圖片從3通道變為64 –> 128 –> 256 –> 512 …. 這些都是卷積層，而我們的視訊記憶體也主要是他們佔用了。
還有上面右邊的params，這些是神經網路的權重大小，可以看到第一層卷積是3×3，而輸入影象的通道是3，輸出通道是64，所以很顯然，第一個卷積層權重所佔的空間是 (3 x 3 x 3) x 64。
另外還有一個需要注意的是中間變數在backward的時候會翻倍！
為什麼，舉個例子，下面是一個計算圖，輸入x，經過中間結果z，然後得到最終變數L：
我們在backward的時候需要儲存下來的中間值。輸出是L，然後輸入x，我們在backward的時候要求L對x的梯度，這個時候就需要在計算鏈L和x中間的z：
dz/dx這個中間值當然要保留下來以用於計算，所以粗略估計，backward的時候中間變數的佔用了是forward的兩倍！
優化器和動量
要注意，優化器也會佔用我們的視訊記憶體！
為什麼,看這個式子:
上式是典型的SGD隨機下降法的總體公式，權重W在進行更新的時候，會產生儲存中間變數
，也就是在優化的時候，模型中的params引數所佔用的視訊記憶體量會翻倍。
當然這只是SGD優化器，其他複雜的優化器如果在計算時需要的中間變數多的時候，就會佔用更多的記憶體。
模型中哪些層會佔用視訊記憶體
有引數的層即會佔用視訊記憶體的層。我們一般的卷積層都會佔用視訊記憶體，而我們經常使用的啟用層Relu沒有引數就不會佔用了。
佔用視訊記憶體的層一般是：
- 卷積層，通常的conv2d
- 全連線層，也就是Linear層
- BatchNorm層
- Embedding層
而不佔用視訊記憶體的則是：
- 剛才說到的啟用層Relu等
- 池化層
- Dropout層
具體計算方式：
- Conv2d(Cin, Cout, K): 引數數目：Cin × Cout × K × K
- Linear(M->N): 引數數目：M×N
- BatchNorm(N): 引數數目： 2N
- Embedding(N,W): 引數數目： N × W
額外的視訊記憶體
總結一下，我們在總體的訓練中，佔用視訊記憶體大概分以下幾類：
- 模型中的引數(卷積層或其他有引數的層)
- 模型在計算時產生的中間引數(也就是輸入影象在計算時每一層產生的輸入和輸出)
- backward的時候產生的額外的中間引數
- 優化器在優化時產生的額外的模型引數
但其實，我們佔用的視訊記憶體空間為什麼比我們理論計算的還要大，原因大概是因為深度學習框架一些額外的開銷吧，不過如果通過上面公式，理論計算出來的視訊記憶體和實際不會差太多的。
如何優化
優化除了演算法層的優化，最基本的優化無非也就一下幾點：
- 減少輸入影象的尺寸
- 減少batch，減少每次的輸入影象數量
- 多使用下采樣，池化層
- 一些神經網路層可以進行小優化，利用relu層中設定
inplace - 購買視訊記憶體更大的顯示卡
- 從深度學習框架上面進行優化
下篇文章我會說明如何在Pytorch這個深度學習框架中跟蹤視訊記憶體的使用量，然後針對Pytorch這個框架進行有目的視訊記憶體優化。
參考：
https://blog.csdn.net/liusandian/article/details/79069926
原文連結：https://ptorch.com/news/181.html
前言
在上篇文章《淺談深度學習:如何計算模型以及中間變數的視訊記憶體佔用大小》中我們對如何計算各種變數所佔視訊記憶體大小進行了一些探索。而這篇文章我們著重講解如何利用Pytorch深度學習框架的一些特性，去檢視我們當前使用的變數所佔用的視訊記憶體大小，以及一些優化工作。以下程式碼所使用的平臺框架為Pytorch。
優化視訊記憶體
在Pytorch中優化視訊記憶體是我們處理大量資料時必要的做法，因為我們並不可能擁有無限的視訊記憶體。視訊記憶體是有限的，而資料是無限的，我們只有優化視訊記憶體的使用量才能夠最大化地利用我們的資料，實現多種多樣的演算法。
估測模型所佔的記憶體
上篇文章中說過，一個模型所佔的視訊記憶體無非是這兩種：
- 模型權重引數
- 模型所儲存的中間變數
其實權重引數一般來說並不會佔用很多的視訊記憶體空間，主要佔用視訊記憶體空間的還是計算時產生的中間變數，當我們定義了一個model之後，我們可以通過以下程式碼簡單計算出這個模型權重引數所佔用的資料量：
import numpy as np
# model是我們在pytorch定義的神經網路層
# model.parameters()取出這個model所有的權重引數
para = sum([np.prod(list(p.size())) for p in model.parameters()])
假設我們有這樣一個model：
Sequential(
(conv_1): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(relu_1): ReLU(inplace)
(conv_2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(relu_2): ReLU(inplace)
(pool_2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
(conv_3): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
)然後我們得到的para是112576，但是我們計算出來的僅僅是權重引數的“數量”，單位是B，我們需要轉化一下：
# 下面的type_size是4，因為我們的引數是float32也就是4B，4個位元組
print('Model {} : params: {:4f}M'.format(model._get_name(), para * type_size / 1000 / 1000))這樣就可以打印出：
Model Sequential : params: 0.450304M但是我們之前說過一個神經網路的模型，不僅僅有權重引數還要計算中間變數的大小。怎麼去計算，我們可以假設一個輸入變數，然後將這個輸入變數投入這個模型中，然後我們主動提取這些計算出來的中間變數：
# model是我們載入的模型
# input是實際中投入的input（Tensor）變數
# 利用clone()去複製一個input，這樣不會對input造成影響
input_ = input.clone()
# 確保不需要計算梯度，因為我們的目的只是為了計算中間變數而已
input_.requires_grad_(requires_grad=False)
mods = list(model.modules())
out_sizes = []
for i in range(1, len(mods)):
m = mods[i]
# 注意這裡，如果relu啟用函式是inplace則不用計算
if isinstance(m, nn.ReLU):
if m.inplace:
continue
out = m(input_)
out_sizes.append(np.array(out.size()))
input_ = out
total_nums = 0
for i in range(len(out_sizes)):
s = out_sizes[i]
nums = np.prod(np.array(s))
total_nums += nums
上面得到的值是模型在執行時候產生所有的中間變數的“數量”，當然我們需要換算一下：
# 列印兩種，只有 forward 和 foreward、backward的情況
print('Model {} : intermedite variables: {:3f} M (without backward)'
.format(model._get_name(), total_nums * type_size / 1000 / 1000))
print('Model {} : intermedite variables: {:3f} M (with backward)'
.format(model._get_name(), total_nums * type_size*2 / 1000 / 1000))
因為在backward的時候所有的中間變數需要儲存下來再來進行計算，所以我們在計算backward的時候，計算出來的中間變數需要乘個2。
然後我們得出，上面這個模型的中間變數需要的佔用的視訊記憶體，很顯然，中間變數佔用的值比模型本身的權重值多多了。如果進行一次backward那麼需要的就更多。
Model Sequential : intermedite variables: 336.089600 M (without backward)
Model Sequential : intermedite variables: 672.179200 M (with backward)我們總結一下之前的程式碼：
# 模型視訊記憶體佔用監測函式
# model：輸入的模型
# input：實際中需要輸入的Tensor變數
# type_size 預設為 4 預設型別為 float32
def modelsize(model, input, type_size=4):
para = sum([np.prod(list(p.size())) for p in model.parameters()])
print('Model {} : params: {:4f}M'.format(model._get_name(), para * type_size / 1000 / 1000))
input_ = input.clone()
input_.requires_grad_(requires_grad=<span class="hljs-keyword">False</span>)
mods = list(model.modules())
out_sizes = []
<span class="hljs-keyword">for</span> i <span class="hljs-keyword">in</span> range(<span class="hljs-number">1</span>, len(mods)):
m = mods[i]
<span class="hljs-keyword">if</span> isinstance(m, nn.ReLU):
<span class="hljs-keyword">if</span> m.inplace:
<span class="hljs-keyword">continue</span>
out = m(input_)
out_sizes.append(np.array(out.size()))
input_ = out
total_nums = <span class="hljs-number">0</span>
<span class="hljs-keyword">for</span> i <span class="hljs-keyword">in</span> range(len(out_sizes)):
s = out_sizes[i]
nums = np.prod(np.array(s))
total_nums += nums
print(<span class="hljs-string">'Model {} : intermedite variables: {:3f} M (without backward)'</span>
.format(model._get_name(), total_nums * type_size / <span class="hljs-number">1000</span> / <span class="hljs-number">1000</span>))
print(<span class="hljs-string">'Model {} : intermedite variables: {:3f} M (with backward)'</span>
.format(model._get_name(), total_nums * type_size*<span class="hljs-number">2</span> / <span class="hljs-number">1000</span> / <span class="hljs-number">1000</span>))
當然我們計算出來的佔用視訊記憶體值僅僅是做參考作用，因為Pytorch在執行的時候需要額外的視訊記憶體值開銷，所以實際的視訊記憶體會比我們計算的稍微大一些。
關於inplace=False
我們都知道啟用函式Relu()有一個預設引數inplace，預設設定為False，當設定為True時，我們在通過relu()計算時的得到的新值不會佔用新的空間而是直接覆蓋原來的值，這也就是為什麼當inplace引數設定為True時可以節省一部分記憶體的緣故。
犧牲計算速度減少視訊記憶體使用量
在Pytorch-0.4.0出來了一個新的功能，可以將一個計算過程分成兩半，也就是如果一個模型需要佔用的視訊記憶體太大了，我們就可以先計算一半，儲存後一半需要的中間結果，然後再計算後一半。
也就是說，新的checkpoint允許我們只儲存反向傳播所需要的部分內容。如果當中缺少一個輸出(為了節省記憶體而導致的)，checkpoint將會從最近的檢查點重新計算中間輸出，以便減少記憶體使用(當然計算時間增加了)：
# 輸入
input = torch.rand(1, 10)
# 假設我們有一個非常深的網路
layers = [nn.Linear(10, 10) for _ in range(1000)]
model = nn.Sequential(*layers)
output = model(input)上面的模型需要佔用很多的記憶體，因為計算中會產生很多的中間變數。為此checkpoint就可以幫助我們來節省記憶體的佔用了。
# 首先設定輸入的input=>requires_grad=True
# 如果不設定可能會導致得到的gradient為0
input = torch.rand(1, 10, requires_grad=True)
layers = [nn.Linear(10, 10) for _ in range(1000)]
# 定義要計算的層函式，可以看到我們定義了兩個
# 一個計算前500個層，另一個計算後500個層
def run_first_half(*args):
x = args[0]
for layer in layers[:500]:
x = layer(x)
return x
def run_second_half(*args):
x = args[0]
for layer in layers[500:-1]:
x = layer(x)
return x
# 我們引入新加的checkpoint
from torch.utils.checkpoint import checkpoint
x = checkpoint(run_first_half, input)
x = checkpoint(run_second_half, x)
# 最後一層單獨調出來執行
x = layers-1
x.sum.backward() # 這樣就可以了
對於Sequential-model來說，因為Sequential()中可以包含很多的block，所以官方提供了另一個功能包：
input = torch.rand(1, 10, requires_grad=True)
layers = [nn.Linear(10, 10) for _ in range(1000)]
model = nn.Sequential(*layers)
from torch.utils.checkpoint import checkpoint_sequential
# 分成兩個部分
num_segments = 2
x = checkpoint_sequential(model, num_segments, input)
x.sum().backward() # 這樣就可以了
跟蹤視訊記憶體使用情況
視訊記憶體的使用情況，在編寫程式中我們可能無法精確計算，但是我們可以通過pynvml這個Nvidia的Python環境庫和Python的垃圾回收工具，可以實時地列印我們使用的視訊記憶體以及哪些Tensor使用了我們的視訊記憶體。
類似於下面的報告：
# 08-Jun-18-17:56:51-gpu_mem_prof
At __main_<span class="hljs-number">_</span> &<span class="hljs-keyword">lt</span>;module&<span class="hljs-keyword">gt</span>;: line <span class="hljs-number">39</span> Total Used Memory:<span class="hljs-number">399.4</span> Mb
At __main_<span class="hljs-number">_</span> &<span class="hljs-keyword">lt</span>;module&<span class="hljs-keyword">gt</span>;: line <span class="hljs-number">40</span> Total Used Memory:<span class="hljs-number">992.5</span> Mb
+ __main_<span class="hljs-number">_</span> &<span class="hljs-keyword">lt</span>;module&<span class="hljs-keyword">gt</span>;: line <span class="hljs-number">40</span> (<span class="hljs-number">1</span>, <span class="hljs-number">1</span>, <span class="hljs-number">682</span>, <span class="hljs-number">700</span>) <span class="hljs-number">1.82</span> M &<span class="hljs-keyword">lt</span>;class <span class="hljs-string">'torch.Tensor'</span>&<span class="hljs-keyword">gt</span>;
+ __main_<span class="hljs-number">_</span> &<span class="hljs-keyword">lt</span>;module&<span class="hljs-keyword">gt</span>;: line <span class="hljs-number">40</span> (<span class="hljs-number">1</span>, <span class="hljs-number">3</span>, <span class="hljs-number">682</span>, <span class="hljs-number">700</span>) <span class="hljs-number">5.46</span> M &<span class="hljs-keyword">lt</span>;class <span class="hljs-string">'torch.Tensor'</span>&<span class="hljs-keyword">gt</span>;
At __main_<span class="hljs-number">_</span> &<span class="hljs-keyword">lt</span>;module&<span class="hljs-keyword">gt</span>;: line <span class="hljs-number">126</span> Total Used Memory:<span class="hljs-number">1088.5</span> Mb
+ __main_<span class="hljs-number">_</span> &<span class="hljs-keyword">lt</span>;module&<span class="hljs-keyword">gt</span>;: line <span class="hljs-number">126</span> (<span class="hljs-number">64</span>, <span class="hljs-number">64</span>, <span class="hljs-number">3</span>, <span class="hljs-number">3</span>) <span class="hljs-number">0</span>.<span class="hljs-number">14</span> M &<span class="hljs-keyword">lt</span>;class <span class="hljs-string">'torch.nn.parameter.Parameter'</span>&<span class="hljs-keyword">gt</span>;
+ __main_<span class="hljs-number">_</span> &<span class="hljs-keyword">lt</span>;module&<span class="hljs-keyword">gt</span>;: line <span class="hljs-number">126</span> (<span class="hljs-number">128</span>, <span class="hljs-number">64</span>, <span class="hljs-number">3</span>, <span class="hljs-number">3</span>) <span class="hljs-number">0</span>.<span class="hljs-number">28</span> M &<span class="hljs-keyword">lt</span>;class <span class="hljs-string">'torch.nn.parameter.Parameter'</span>&<span class="hljs-keyword">gt</span>;
+ __main_<span class="hljs-number">_</span> &<span class="hljs-keyword">lt</span>;module&<span class="hljs-keyword">gt</span>;: line <span class="hljs-number">126</span> (<span class="hljs-number">128</span>, <span class="hljs-number">128</span>, <span class="hljs-number">3</span>, <span class="hljs-number">3</span>) <span class="hljs-number">0</span>.<span class="hljs-number">56</span> M &<span class="hljs-keyword">lt</span>;class <span class="hljs-string">'torch.nn.parameter.Parameter'</span>&<span class="hljs-keyword">gt</span>;
+ __main_<span class="hljs-number">_</span> &<span class="hljs-keyword">lt</span>;module&<span class="hljs-keyword">gt</span>;: line <span class="hljs-number">126</span> (<span class="hljs-number">64</span>, <span class="hljs-number">3</span>, <span class="hljs-number">3</span>, <span class="hljs-number">3</span>) <span class="hljs-number">0</span>.<span class="hljs-number">00</span> M &<span class="hljs-keyword">lt</span>;class <span class="hljs-string">'torch.nn.parameter.Parameter'</span>&<span class="hljs-keyword">gt</span>;
+ __main_<span class="hljs-number">_</span> &<span class="hljs-keyword">lt</span>;module&<span class="hljs-keyword">gt</span>;: line <span class="hljs-number">126</span> (<span class="hljs-number">256</span>, <span class="hljs-number">256</span>, <span class="hljs-number">3</span>, <span class="hljs-number">3</span>) <span class="hljs-number">2.25</span> M &<span class="hljs-keyword">lt</span>;class <span class="hljs-string">'torch.nn.parameter.Parameter'</span>&<span class="hljs-keyword">gt</span>;
+ __main_<span class="hljs-number">_</span> &<span class="hljs-keyword">lt</span>;module&<span class="hljs-keyword">gt</span>;: line <span class="hljs-number">126</span> (<span class="hljs-number">512</span>, <span class="hljs-number">256</span>, <span class="hljs-number">3</span>, <span class="hljs-number">3</span>) <span class="hljs-number">4.5</span> M &<span class="hljs-keyword">lt</span>;class <span class="hljs-string">'torch.nn.parameter.Parameter'</span>&<span class="hljs-keyword">gt</span>;
+ __main_<span class="hljs-number">_</span> &<span class="hljs-keyword">lt</span>;module&<span class="hljs-keyword">gt</span>;: line <span class="hljs-number">126</span> (<span class="hljs-number">512</span>, <span class="hljs-number">512</span>, <span class="hljs-number">3</span>, <span class="hljs-number">3</span>) <span class="hljs-number">9.0</span> M &<span class="hljs-keyword">lt</span>;class <span class="hljs-string">'torch.nn.parameter.Parameter'</span>&<span class="hljs-keyword">gt</span>;
+ __main_<span class="hljs-number">_</span> &<span class="hljs-keyword">lt</span>;module&<span class="hljs-keyword">gt</span>;: line <span class="hljs-number">126</span> (<span class="hljs-number">64</span>,) <span class="hljs-number">0</span>.<span class="hljs-number">00</span> M &<span class="hljs-keyword">lt</span>;class <span class="hljs-string">'torch.nn.parameter.Parameter'</span>&<span class="hljs-keyword">gt</span>;
+ __main_<span class="hljs-number">_</span> &<span class="hljs-keyword">lt</span>;module&<span class="hljs-keyword">gt</span>;: line <span class="hljs-number">126</span> (<span class="hljs-number">1</span>, <span class="hljs-number">3</span>, <span class="hljs-number">682</span>, <span class="hljs-number">700</span>) <span class="hljs-number">5.46</span> M &<span class="hljs-keyword">lt</span>;class <span class="hljs-string">'torch.Tensor'</span>&<span class="hljs-keyword">gt</span>;
+ __main_<span class="hljs-number">_</span> &<span class="hljs-keyword">lt</span>;module&<span class="hljs-keyword">gt</span>;: line <span class="hljs-number">126</span> (<span class="hljs-number">128</span>,) <span class="hljs-number">0</span>.<span class="hljs-number">00</span> M &<span class="hljs-keyword">lt</span>;class <span class="hljs-string">'torch.nn.parameter.Parameter'</span>&<span class="hljs-keyword">gt</span>;
+ __main_<span class="hljs-number">_</span> &<span class="hljs-keyword">lt</span>;module&<span class="hljs-keyword">gt</span>;: line <span class="hljs-number">126</span> (<span class="hljs-number">256</span>,) <span class="hljs-number">0</span>.<span class="hljs-number">00</span> M &<span class="hljs-keyword">lt</span>;class <span class="hljs-string">'torch.nn.parameter.Parameter'</span>&<span class="hljs-keyword">gt</span>;
+ __main_<span class="hljs-number">_</span> &<span class="hljs-keyword">lt</span>;module&<span class="hljs-keyword">gt</span>;: line <span class="hljs-number">126</span> (<span class="hljs-number">512</span>,) <span class="hljs-number">0</span>.<span class="hljs-number">00</span> M &<span class="hljs-keyword">lt</span>;class <span class="hljs-string">'torch.nn.parameter.Parameter'</span>&<span class="hljs-keyword">gt</span>;
+ __main_<span class="hljs-number">_</span> &<span class="hljs-keyword">lt</span>;module&<span class="hljs-keyword">gt</span>;: line <span class="hljs-number">126</span> (<span class="hljs-number">3</span>,) <span class="hljs-number">1.14</span> M &<span class="hljs-keyword">lt</span>;class <span class="hljs-string">'torch.Tensor'</span>&<span class="hljs-keyword">gt</span>;
+ __main_<span class="hljs-number">_</span> &<span class="hljs-keyword">lt</span>;module&<span class="hljs-keyword">gt</span>;: line <span class="hljs-number">126</span> (<span class="hljs-number">256</span>, <span class="hljs-number">128</span>, <span class="hljs-number">3</span>, <span class="hljs-number">3</span>) <span class="hljs-number">1.12</span> M &<span class="hljs-keyword">lt</span>;class <span class="hljs-string">'torch.nn.parameter.Parameter'</span>&<span class="hljs-keyword">gt</span>;
...</code></pre>
以下是相關的程式碼，目前程式碼依然有些地方需要修改，等修改完善好我會將完整程式碼以及使用說明放到github上：https://github.com/Oldpan/Pytorch-Memory-Utils
請大家多多留意。
import datetime
import linecache
import os
import gc
import pynvml
import torch
import numpy as np
print_tensor_sizes = True
last_tensor_sizes = set()
gpu_profile_fn = f'{datetime.datetime.now():%d-%b-%y-%H:%M:%S}-gpu_mem_prof.txt'
# if 'GPU_DEBUG' in os.environ:
# print('profiling gpu usage to ', gpu_profile_fn)
lineno = None
func_name = None
filename = None
module_name = None
# fram = inspect.currentframe()
# func_name = fram.f_code.co_name
# filename = fram.f_globals["__file__"]
# ss = os.path.dirname(os.path.abspath(filename))
# module_name = fram.f_globals["__name__"]
def gpu_profile(frame, event):
# it is _about to_ execute (!)
global last_tensor_sizes
global lineno, func_name, filename, module_name
if event == 'line':
try:
# about _previous_ line (!)
if lineno is not None:
pynvml.nvmlInit()
# handle = pynvml.nvmlDeviceGetHandleByIndex(int(os.environ['GPU_DEBUG']))
handle = pynvml.nvmlDeviceGetHandleByIndex(0)
meminfo = pynvml.nvmlDeviceGetMemoryInfo(handle)
line = linecache.getline(filename, lineno)
where_str = module_name+' '+func_name+':'+' line '+str(lineno)
with open(gpu_profile_fn, 'a+') as f:
f.write(f"At {where_str:<50}"
f"Total Used Memory:{meminfo.used/1024**2:<7.1f}Mb\n")
if print_tensor_sizes is True:
for tensor in get_tensors():
if not hasattr(tensor, 'dbg_alloc_where'):
tensor.dbg_alloc_where = where_str
new_tensor_sizes = {(type(x), tuple(x.size()), np.prod(np.array(x.size()))*4/1024**2,
x.dbg_alloc_where) for x in get_tensors()}
for t, s, m, loc in new_tensor_sizes - last_tensor_sizes:
f.write(f'+ {loc:<50} {str(s):<20} {str(m)[:4]} M {str(t):<10}\n')
for t, s, m, loc in last_tensor_sizes - new_tensor_sizes:
f.write(f'- {loc:<50} {str(s):<20} {str(m)[:4]} M {str(t):<10}\n')
last_tensor_sizes = new_tensor_sizes
pynvml.nvmlShutdown()
# save details about line _to be_ executed
lineno = None
func_name = frame.f_code.co_name
filename = frame.f_globals["__file__"]
if (filename.endswith(".pyc") or
filename.endswith(".pyo")):
filename = filename[:-1]
module_name = frame.f_globals["__name__"]
lineno = frame.f_lineno
return gpu_profile
except Exception as e:
print('A exception occured: {}'.format(e))
return gpu_profile
def get_tensors():
for obj in gc.get_objects():
try:
if torch.is_tensor(obj):
tensor = obj
else:
continue
if tensor.is_cuda:
yield tensor
except Exception as e:
print('A exception occured: {}'.format(e))
需要注意的是，linecache中的getlines只能讀取緩衝過的檔案，如果這個檔案沒有執行過則返回無效值。Python 的垃圾收集機制會在變數沒有應引用的時候立馬進行回收，但是為什麼模型中計算的中間變數在執行結束後還會存在呢。既然都沒有引用了為什麼還會佔用空間？
一種可能的情況是這些引用不在Python程式碼中，而是在神經網路層的執行中為了backward被儲存為gradient，這些引用都在計算圖中，我們在程式中是無法看到的：
後記
實際中我們會有些只使用一次的模型，為了節省視訊記憶體，我們需要一邊計算一遍清除中間變數，使用del進行操作。限於篇幅這裡不進行講解，下一篇會進行說明。
<br>
原創文章，轉載請註明 ：<a href="https://ptorch.com/news/181.html" target="_blank">如何在Pytorch中精細化利用視訊記憶體以及提高Pytorch視訊記憶體利用率 - pytorch中文網</a><br>
原文出處： https://ptorch.com/news/181.html<br>
問題交流群 ：168117787
</div>