關於python中顯示記憶體回收的問題
技術背景
筆者在執行一個Jax的任務中，又發現了一個奇怪的問題，就是明明只分配了很小的矩陣空間，但是在多次的任務執行之後，顯示記憶體突然就爆了。而且此時已經按照Jax的官方說明配置了XLA_PYTHON_CLIENT_PREALLOCATE這個參數為false，也就是不進行顯示記憶體的預分配（默認會分配90%的顯示記憶體空間以供使用）。然後在網上找到了一些類似的問題，比如參考鏈接中的1、2、3、4，都是在一些操作後發現未釋放顯示記憶體，這裡提供一個實例問題和處理的思路，如果有更好的方案歡迎大家在評論區留言。
問題復現
在未執行任何GPU的任務時，我們可以看到此時nvidia-smi的輸出如下：
Tue Dec 14 16:14:32 2021
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 470.42.01 Driver Version: 470.42.01 CUDA Version: 11.4 |
|-------------------------------+----------------------+----------------------+
| GPU Name Persistence-M| Bus-Id Disp.A | Volatile Uncorr. ECC |
| Fan Temp Perf Pwr:Usage/Cap| Memory-Usage | GPU-Util Compute M. |
| | | MIG M. |
|===============================+======================+======================|
| 0 Quadro RTX 4000 On | 00000000:03:00.0 On | N/A |
| 30% 43C P8 20W / 125W | 1260MiB / 7979MiB | 10% Default |
| | | N/A |
+-------------------------------+----------------------+----------------------+
| 1 Quadro RTX 4000 On | 00000000:A6:00.0 Off | N/A |
| 30% 34C P8 7W / 125W | 10MiB / 7982MiB | 0% Default |
| | | N/A |
+-------------------------------+----------------------+----------------------+
+-----------------------------------------------------------------------------+
| Processes: |
| GPU GI CI PID Type Process name GPU Memory |
| ID ID Usage |
|=============================================================================|
| 0 N/A N/A 1673 G /usr/lib/xorg/Xorg 110MiB |
| 0 N/A N/A 3015 G /usr/lib/xorg/Xorg 661MiB |
| 0 N/A N/A 3251 G /usr/bin/gnome-shell 132MiB |
| 0 N/A N/A 1142734 G ...AAAAAAAAA= --shared-files 64MiB |
| 0 N/A N/A 1337710 G ...AAAAAAAAA= --shared-files 80MiB |
| 0 N/A N/A 1371509 G ...369783.log --shared-files 63MiB |
| 0 N/A N/A 1506625 G ...AAAAAAAAA= --shared-files 89MiB |
| 1 N/A N/A 1673 G /usr/lib/xorg/Xorg 4MiB |
| 1 N/A N/A 3015 G /usr/lib/xorg/Xorg 4MiB |
+-----------------------------------------------------------------------------+
此時啟動一個ipython的終端窗口，執行如下的Jax任務：
In [1]: import numpy as np
In [2]: import os
...: os.environ['CUDA_VISIBLE_DEVICES']='1'
...: os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false"
In [3]: from jax import numpy as jnp
In [4]: a = np.ones(1000000)
In [5]: b = jnp.array(a)
此時再次查看nvidia-smi的結果如下：
Tue Dec 14 16:18:26 2021
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 470.42.01 Driver Version: 470.42.01 CUDA Version: 11.4 |
|-------------------------------+----------------------+----------------------+
| GPU Name Persistence-M| Bus-Id Disp.A | Volatile Uncorr. ECC |
| Fan Temp Perf Pwr:Usage/Cap| Memory-Usage | GPU-Util Compute M. |
| | | MIG M. |
|===============================+======================+======================|
| 0 Quadro RTX 4000 On | 00000000:03:00.0 On | N/A |
| 30% 42C P8 20W / 125W | 1238MiB / 7979MiB | 10% Default |
| | | N/A |
+-------------------------------+----------------------+----------------------+
| 1 Quadro RTX 4000 On | 00000000:A6:00.0 Off | N/A |
| 30% 36C P0 35W / 125W | 114MiB / 7982MiB | 0% Default |
| | | N/A |
+-------------------------------+----------------------+----------------------+
+-----------------------------------------------------------------------------+
| Processes: |
| GPU GI CI PID Type Process name GPU Memory |
| ID ID Usage |
|=============================================================================|
| 0 N/A N/A 1673 G /usr/lib/xorg/Xorg 110MiB |
| 0 N/A N/A 3015 G /usr/lib/xorg/Xorg 661MiB |
| 0 N/A N/A 3251 G /usr/bin/gnome-shell 129MiB |
| 0 N/A N/A 1142734 G ...AAAAAAAAA= --shared-files 44MiB |
| 0 N/A N/A 1337710 G ...AAAAAAAAA= --shared-files 80MiB |
| 0 N/A N/A 1371509 G ...369783.log --shared-files 63MiB |
| 0 N/A N/A 1506625 G ...AAAAAAAAA= --shared-files 89MiB |
| 1 N/A N/A 1673 G /usr/lib/xorg/Xorg 4MiB |
| 1 N/A N/A 3015 G /usr/lib/xorg/Xorg 4MiB |
| 1 N/A N/A 1743467 C /usr/local/bin/python 101MiB |
+-----------------------------------------------------------------------------+
此時的結果還是比較符合我們的預期的，這個python的進程佔用了101MB的空間。但是此時如果我們在ipython中把這個對象刪除了：
In [6]: del b
In [7]: b
---------------------------------------------------------------------------
NameError Traceback (most recent call last)
<ipython-input-7-89e6c98d9288> in <module>
----> 1 b
NameError: name 'b' is not defined
然後再次查看nvidia-smi的結果：
Tue Dec 14 16:21:12 2021
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 470.42.01 Driver Version: 470.42.01 CUDA Version: 11.4 |
|-------------------------------+----------------------+----------------------+
| GPU Name Persistence-M| Bus-Id Disp.A | Volatile Uncorr. ECC |
| Fan Temp Perf Pwr:Usage/Cap| Memory-Usage | GPU-Util Compute M. |
| | | MIG M. |
|===============================+======================+======================|
| 0 Quadro RTX 4000 On | 00000000:03:00.0 On | N/A |
| 30% 42C P5 21W / 125W | 1231MiB / 7979MiB | 0% Default |
| | | N/A |
+-------------------------------+----------------------+----------------------+
| 1 Quadro RTX 4000 On | 00000000:A6:00.0 Off | N/A |
| 30% 35C P8 7W / 125W | 114MiB / 7982MiB | 0% Default |
| | | N/A |
+-------------------------------+----------------------+----------------------+
+-----------------------------------------------------------------------------+
| Processes: |
| GPU GI CI PID Type Process name GPU Memory |
| ID ID Usage |
|=============================================================================|
| 0 N/A N/A 1673 G /usr/lib/xorg/Xorg 110MiB |
| 0 N/A N/A 3015 G /usr/lib/xorg/Xorg 662MiB |
| 0 N/A N/A 3251 G /usr/bin/gnome-shell 111MiB |
| 0 N/A N/A 1142734 G ...AAAAAAAAA= --shared-files 55MiB |
| 0 N/A N/A 1337710 G ...AAAAAAAAA= --shared-files 80MiB |
| 0 N/A N/A 1371509 G ...369783.log --shared-files 63MiB |
| 0 N/A N/A 1506625 G ...AAAAAAAAA= --shared-files 89MiB |
| 1 N/A N/A 1673 G /usr/lib/xorg/Xorg 4MiB |
| 1 N/A N/A 3015 G /usr/lib/xorg/Xorg 4MiB |
| 1 N/A N/A 1743467 C /usr/local/bin/python 101MiB |
+-----------------------------------------------------------------------------+
此時我們可以看到，雖然已經把對象給刪除了，在python的程式中已然找不到這個對象，但是在顯示記憶體中的數據並未被消除。而且如果一直不消除，這塊顯示記憶體就會一直佔用在那裡，直到python進程（此時作為該進程的一個守護進程）的結束。
解決思路
暫時還不清楚這個問題發生的機制，在一些特定場景下出現殭屍進程的問題似乎跟我復現的這個場景也有所不同。只是考慮到在python的進程結束之後，這一塊的顯示記憶體還是被成功釋放了的，因此我考慮直接用進程的方法來解決這個顯示記憶體分配和清空的方法，以下是一個基於進程實現的案例：
import os
os.environ['CUDA_VISIBLE_DEVICES']='1'
os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false"
import time
from multiprocessing import Pool
import numpy as np
from jax import numpy as jnp
a = np.ones(1000000)
def f(a):
b = jnp.array(a)
time.sleep(2)
print('Array b has been deleted!')
return True
with Pool(1) as p:
res = p.map(f, [(a,)])
print ('Is jax array deleted successfully?\t{}'.format(res))
time.sleep(6)
在這個程式中，我們把要執行的相關任務，包含GPU矩陣的轉化與分配，都放到了一個函數中，然後通過multiprocessing開啟一個子進程，來執行這個任務，並且在任務中甚至不需要手動執行del這個刪除的操作。這麼一來，我們既可以實現對象的即時銷毀，也通過進程式控制制的機制確保在顯示記憶體中佔用的位置被清空。如果進程執行中存在一些問題，還可以通過terminate的操作來直接殺死進程，同樣也可以確保顯示記憶體佔用不會發生堆積的情況。程式的執行結果如下：
Array b has been deleted!
Is jax array deleted successfully? [True]
在程式執行的過程中我們也可以看到，在nvidia-smi中的顯示記憶體佔用，僅僅持續了2秒，也就是我們在函數內部設置的進程sleep參數。而在之後6秒的sleep時間中，這一塊記憶體佔用是被清空了的，這也就達到了我們最初的目的。當然，最重要的是，我們依然可以從函數中獲取到返回值，這就確保後面有需要存儲或者使用到的參數不被同步的銷毀。需要注意的是，在同等條件下，如果不使用子進程來執行這個函數，而是直接使用res=f(a)的形式來執行，作為臨時變數的b最終依然存在於顯示記憶體之中，這是一個非常可怕的事情。
總結概要
在使用一些python的GPU模組，或者寫CUDA時，有時會發現顯示記憶體被無端佔用的場景，即時執行了cudaFree()或者python的del操作，也無法消除這一塊的顯示記憶體佔用。最終我們發現，可以通過額外開啟一個子進程的方法來封裝相關的操作，通過對進程的存活控制來實現對GPU顯示記憶體佔用的控制，有可能是一個臨時規避問題的思路。
版權聲明
本文首發鏈接為：//www.cnblogs.com/dechinphy/p/gc.html
作者ID：DechinPhy
更多原著文章請參考：//www.cnblogs.com/dechinphy/
打賞專用鏈接：//www.cnblogs.com/dechinphy/gallery/image/379634.html
騰訊雲專欄同步：//cloud.tencent.com/developer/column/91958
參考鏈接
- //blog.csdn.net/jzrita/article/details/80719297
- //blog.csdn.net/xxs8457800/article/details/104307283
- //jermine.vdo.pub/python/解決gpu顯示記憶體未釋放問題/
- //blog.csdn.net/weixin_42317730/article/details/116786526?share_token=7ef0f7d6-6d68-4efb-995b-24517000ac11&tt_from=copy_link&utm_source=copy_link&utm_medium=toutiao_android&utm_campaign=client_share?=linux清理gpu記憶體,GPU記憶體在CUDA腳本執行後無法
