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