python中顯存回收問題解決方法
1.技術背景
筆者在執行一個Jax的任務中，又發現瞭一個奇怪的問題，就是明明隻分配瞭很小的矩陣空間，但是在多次的任務執行之後，顯存突然就爆瞭。而且此時已經按照Jax的官方說明配置瞭XLA_PYTHON_CLIENT_PREALLOCATE這個參數為false，也就是不進行顯存的預分配（默認會分配90%的顯存空間以供使用）。然後在網上找到瞭一些類似的問題，比如參考鏈接中的1、2、3、4，都是在一些操作後發現未釋放顯存，這裡提供一個實例問題和處理的思路，如果有更好的方案歡迎大傢在評論區留言。
2.問題復現
在未執行任何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進程（此時作為該進程的一個守護進程）的結束。
3.解決思路
暫時還不清楚這個問題發生的機制，在一些特定場景下出現僵屍進程的問題似乎跟我復現的這個場景也有所不同。隻是考慮到在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最終依然存在於顯存之中，這是一個非常可怕的事情。
4.總結概要
在使用一些python的GPU模塊，或者寫CUDA時，有時會發現顯存被無端占用的場景，即時執行瞭cudaFree()或者python的del操作，也無法消除這一塊的顯存占用。最終我們發現，可以通過額外開啟一個子進程的方法來封裝相關的操作，通過對進程的存活控制來實現對GPU顯存占用的控制，有可能是一個臨時規避問題的思路。
到此這篇關於python中顯存回收問題解決方法的文章就介紹到這瞭,更多相關python顯存回收內容請搜索WalkonNet以前的文章或繼續瀏覽下面的相關文章希望大傢以後多多支持WalkonNet！
推薦閱讀：
- Win10下安裝CUDA11.0+CUDNN8.0+tensorflow-gpu2.4.1+pytorch1.7.0+paddlepaddle-gpu2.0.0
- 關於PyTorch環境配置及安裝教程(Windows10)
- win10+RTX3050ti+TensorFlow+cudn+cudnn配置深度學習環境的方法
- 詳解win10下pytorch-gpu安裝以及CUDA詳細安裝過程
- Python深度學習之Pytorch初步使用