python中顯存回收問題解決方法
1.技術(shù)背景
筆者在執(zhí)行一個(gè)Jax的任務(wù)中,又發(fā)現(xiàn)了一個(gè)奇怪的問題,就是明明只分配了很小的矩陣空間,但是在多次的任務(wù)執(zhí)行之后,顯存突然就爆了。而且此時(shí)已經(jīng)按照J(rèn)ax的官方說明配置了XLA_PYTHON_CLIENT_PREALLOCATE這個(gè)參數(shù)為false,也就是不進(jìn)行顯存的預(yù)分配(默認(rèn)會(huì)分配90%的顯存空間以供使用)。然后在網(wǎng)上找到了一些類似的問題,比如參考鏈接中的1、2、3、4,都是在一些操作后發(fā)現(xiàn)未釋放顯存,這里提供一個(gè)實(shí)例問題和處理的思路,如果有更好的方案歡迎大家在評(píng)論區(qū)留言。
2.問題復(fù)現(xiàn)
在未執(zhí)行任何GPU的任務(wù)時(shí),我們可以看到此時(shí)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 |
+-----------------------------------------------------------------------------+
此時(shí)啟動(dòng)一個(gè)ipython的終端窗口,執(zhí)行如下的Jax任務(wù):
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)
此時(shí)再次查看nvidia-smi的結(jié)果如下:
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 |
+-----------------------------------------------------------------------------+
此時(shí)的結(jié)果還是比較符合我們的預(yù)期的,這個(gè)python的進(jìn)程占用了101MB的空間。但是此時(shí)如果我們?cè)趇python中把這個(gè)對(duì)象刪除了:
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的結(jié)果:
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 |
+-----------------------------------------------------------------------------+
此時(shí)我們可以看到,雖然已經(jīng)把對(duì)象給刪除了,在python的程序中已然找不到這個(gè)對(duì)象,但是在顯存中的數(shù)據(jù)并未被消除。而且如果一直不消除,這塊顯存就會(huì)一直占用在那里,直到python進(jìn)程(此時(shí)作為該進(jìn)程的一個(gè)守護(hù)進(jìn)程)的結(jié)束。
3.解決思路
暫時(shí)還不清楚這個(gè)問題發(fā)生的機(jī)制,在一些特定場景下出現(xiàn)僵尸進(jìn)程的問題似乎跟我復(fù)現(xiàn)的這個(gè)場景也有所不同。只是考慮到在python的進(jìn)程結(jié)束之后,這一塊的顯存還是被成功釋放了的,因此我考慮直接用進(jìn)程的方法來解決這個(gè)顯存分配和清空的方法,以下是一個(gè)基于進(jìn)程實(shí)現(xiàn)的案例:
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)
在這個(gè)程序中,我們把要執(zhí)行的相關(guān)任務(wù),包含GPU矩陣的轉(zhuǎn)化與分配,都放到了一個(gè)函數(shù)中,然后通過multiprocessing開啟一個(gè)子進(jìn)程,來執(zhí)行這個(gè)任務(wù),并且在任務(wù)中甚至不需要手動(dòng)執(zhí)行del這個(gè)刪除的操作。這么一來,我們既可以實(shí)現(xiàn)對(duì)象的即時(shí)銷毀,也通過進(jìn)程控制的機(jī)制確保在顯存中占用的位置被清空。如果進(jìn)程執(zhí)行中存在一些問題,還可以通過terminate的操作來直接殺死進(jìn)程,同樣也可以確保顯存占用不會(huì)發(fā)生堆積的情況。程序的執(zhí)行結(jié)果如下:
Array b has been deleted! Is jax array deleted successfully? [True]
在程序執(zhí)行的過程中我們也可以看到,在nvidia-smi中的顯存占用,僅僅持續(xù)了2秒,也就是我們?cè)诤瘮?shù)內(nèi)部設(shè)置的進(jìn)程sleep參數(shù)。而在之后6秒的sleep時(shí)間中,這一塊內(nèi)存占用是被清空了的,這也就達(dá)到了我們最初的目的。當(dāng)然,最重要的是,我們依然可以從函數(shù)中獲取到返回值,這就確保后面有需要存儲(chǔ)或者使用到的參數(shù)不被同步的銷毀。需要注意的是,在同等條件下,如果不使用子進(jìn)程來執(zhí)行這個(gè)函數(shù),而是直接使用res=f(a)的形式來執(zhí)行,作為臨時(shí)變量的b最終依然存在于顯存之中,這是一個(gè)非??膳碌氖虑?。
4.總結(jié)概要
在使用一些python的GPU模塊,或者寫CUDA時(shí),有時(shí)會(huì)發(fā)現(xiàn)顯存被無端占用的場景,即時(shí)執(zhí)行了cudaFree()或者python的del操作,也無法消除這一塊的顯存占用。最終我們發(fā)現(xiàn),可以通過額外開啟一個(gè)子進(jìn)程的方法來封裝相關(guān)的操作,通過對(duì)進(jìn)程的存活控制來實(shí)現(xiàn)對(duì)GPU顯存占用的控制,有可能是一個(gè)臨時(shí)規(guī)避問題的思路。
到此這篇關(guān)于python中顯存回收問題解決方法的文章就介紹到這了,更多相關(guān)python顯存回收內(nèi)容請(qǐng)搜索腳本之家以前的文章或繼續(xù)瀏覽下面的相關(guān)文章希望大家以后多多支持腳本之家!
相關(guān)文章
Python實(shí)現(xiàn)向PPT中插入表格與圖片的方法詳解
這篇文章將帶大家學(xué)習(xí)一下如何在PPT中插入表格與圖片以及在表格中插入內(nèi)容,文中的示例代碼講解詳細(xì),感興趣的小伙伴可以跟隨小編一起學(xué)習(xí)一下2022-05-05
Python中的類型提示(Type Hints)總結(jié)
Python3.5 版本引入了類型提示(Type Hints),它允許開發(fā)者在代碼中顯式地聲明變量、函數(shù)、方法等的類型信息,下面小編就來帶大家一起看看Python類型提示的初步使用吧2023-05-05
Python利用Beautiful Soup模塊搜索內(nèi)容詳解
這篇文章主要給大家介紹了python中 Beautiful Soup 模塊的搜索方法函數(shù)。 方法不同類型的過濾參數(shù)能夠進(jìn)行不同的過濾,得到想要的結(jié)果。文中介紹的非常詳細(xì),對(duì)大家具有一定的參考價(jià)值,需要的朋友們下面來一起看看吧。2017-03-03
Python中在腳本中引用其他文件函數(shù)的實(shí)現(xiàn)方法
下面小編就為大家?guī)硪黄狿ython中在腳本中引用其他文件函數(shù)的實(shí)現(xiàn)方法。小編覺得挺不錯(cuò)的,現(xiàn)在就分享給大家,也給大家做個(gè)參考。一起跟隨小編過來看看吧2016-06-06
python實(shí)戰(zhàn)之德州撲克第二步-判斷牌型
這篇文章主要介紹了python實(shí)戰(zhàn)之德州撲克第二步-判斷牌型,文中有非常詳細(xì)的代碼示例,對(duì)正在學(xué)習(xí)python的小伙伴們有非常好的幫助,需要的朋友可以參考下2021-04-04

