前言

訓練模型時，一般我們會把模型model，數(shù)據(jù)data和標簽label放到GPU顯存中進行加速。

但有的時候GPU Memory會增加，有的時候會保持不變，以及我們要怎么清理掉一些用完的變量呢？

下面讓我們一起來探究下原理吧！

一、pytorch訓練模型

只要你把任何東西（無論是多小的tensor）放到GPU顯存中，那么你至少會棧1000MiB左右的顯存（根據(jù)cuda版本，會略有不同）。這部分顯存是cuda running時固有配件必須要占掉的顯存，你先訓練過程匯總也是無法釋放的。

import torch

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# 把一個很小的tensor加到GPU顯存中（大約2MiB）
x = torch.randn((2, 3), device=device) 

現(xiàn)在我再放入一個比較大的tensor，GPU顯存升到了1919MiB

y = torch.randn((200, 300, 200, 20), device=device)

也就是說當你有個新的東西加進去時，GPU顯存會不斷擴大。

二、batch訓練模型時，GPU顯存的容量會保持不變?

但是為什么我們像下面這樣拿出一個個batch訓練模型時，GPU顯存的容量會保持不變的呢？

batch_loss = []
for epoch in range(self.epoch):
    pbar = enumerate(tqdm(self.train_dataloader, desc="Training Bar"), 0)
    for i, (inputs, labels) in pbar:
        inputs = inputs.to(self.device) # 放到GPU顯存中
        labels = labels.to(self.device) # 放到GPU顯存中

        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)  
        batch_loss.append(loss)
        loss.backward()
        optimizer.step()

這實際上跟cuda的內(nèi)存管理有關(guān)，但你的一塊內(nèi)存不再由引用它的變量時，他就不再是Active memory。他會殘留在一個隊列中，如果下一次有新的東西進來，那就會把他擠出隊列（FIFO），當然如果新進來的東西和你那部分空間的大小差不多，那么出去一塊空間，又進來一塊空間，那么看上去你的GPU顯存就不會有增加，看上去是一直不變。這就解釋了一個個batch訓練模型時GPU顯存不變的原因。

當然如果新加進來的東西很多，就是那些unactivate的memory全部被擠出還是放不下新的東西，那GPU的顯存就會增加了。（有點類似C++中的capacity增加的情況）

實際運行中，我們會發(fā)現(xiàn)這個隊列capacity會有個閾值（這個閾值好像不是固定的），當你還沒到達這個閾值時不會觸發(fā)垃圾回收機制（即清理unactivate memeory的空間）

也就是說我不斷運行下面代碼

y = torch.randn((200, 300, 200, 20), device=device)

GPU顯存會增加到9247MiB

這時再加入y，容量不再增加，會把原來unactivate的memory擠掉。

三、如何釋放GPU顯存空間

那么我們要怎么樣釋放掉空間呢

我們上面很多空間原來是被y指向的，后來y指向新的地方，那這些空間都是殘留的，我們可以用下面命令繼續(xù)釋放（如果你想釋放的話）

torch.cuda.empty_cache()

上述命令可能要運行多次才會釋放空間，我運行了大概5次吧

殘留內(nèi)存成功被釋放

現(xiàn)在這里面GPU顯存 = 基礎(chǔ)配置(1001MiB) + y(918MiB) + x(忽略不計)

最后我們再來把y這部分釋放掉

令 y = 2，那么原來y所指的那部分顯存空間就會變成unactivate，我們可以使用torch.cuda.empty_cache()把這部分空間釋放掉

最終只剩下基礎(chǔ)配置的GPU顯存占用（這部分已經(jīng)無法釋放了）

四、torch.cuda.memory_summary()查看顯存信息

使用print(torch.cuda.memory_summary())可以看到更多關(guān)于cuda顯存的信息

五、寫在最后

經(jīng)過上面的摸索，我感覺這部分內(nèi)容跟操作系統(tǒng)的內(nèi)存管理有點像，所以說計算機的那幾門基礎(chǔ)課真的很重要，大家都要好好學一學！

當然實際中cuda的顯存管理肯定沒有那么簡單，有興趣的同學可以繼續(xù)探究下。