python中顯存回收問題解決方法

更新時間：2021年12月15日 10:19:13 作者：Dechin's Blog

這篇文章主要介紹了python中顯存回收問題解決方法，文章提供一個實例問題和處理的思路，具有一定的參考價值，需要的小伙伴可以參考一下

1.技術(shù)背景

筆者在執(zhí)行一個Jax的任務(wù)中，又發(fā)現(xiàn)了一個奇怪的問題，就是明明只分配了很小的矩陣空間，但是在多次的任務(wù)執(zhí)行之后，顯存突然就爆了。而且此時已經(jīng)按照J(rèn)ax的官方說明配置了XLA_PYTHON_CLIENT_PREALLOCATE這個參數(shù)為false，也就是不進行顯存的預(yù)分配（默認(rèn)會分配90%的顯存空間以供使用）。然后在網(wǎng)上找到了一些類似的問題，比如參考鏈接中的1、2、3、4，都是在一些操作后發(fā)現(xiàn)未釋放顯存，這里提供一個實例問題和處理的思路，如果有更好的方案歡迎大家在評論區(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的進程占用了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進程（此時作為該進程的一個守護進程）的結(jié)束。

3.解決思路

暫時還不清楚這個問題發(fā)生的機制，在一些特定場景下出現(xiàn)僵尸進程的問題似乎跟我復(fù)現(xiàn)的這個場景也有所不同。只是考慮到在python的進程結(jié)束之后，這一塊的顯存還是被成功釋放了的，因此我考慮直接用進程的方法來解決這個顯存分配和清空的方法，以下是一個基于進程實現(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開啟一個子進程，來執(zhí)行這個任務(wù)，并且在任務(wù)中甚至不需要手動執(zhí)行del這個刪除的操作。這么一來，我們既可以實現(xiàn)對象的即時銷毀，也通過進程控制的機制確保在顯存中占用的位置被清空。如果進程執(zhí)行中存在一些問題，還可以通過terminate的操作來直接殺死進程，同樣也可以確保顯存占用不會發(fā)生堆積的情況。程序的執(zhí)行結(jié)果如下：

Array b has been deleted!
Is jax array deleted successfully?      [True]

在程序執(zhí)行的過程中我們也可以看到，在nvidia-smi中的顯存占用，僅僅持續(xù)了2秒，也就是我們在函數(shù)內(nèi)部設(shè)置的進程sleep參數(shù)。而在之后6秒的sleep時間中，這一塊內(nèi)存占用是被清空了的，這也就達到了我們最初的目的。當(dāng)然，最重要的是，我們依然可以從函數(shù)中獲取到返回值，這就確保后面有需要存儲或者使用到的參數(shù)不被同步的銷毀。需要注意的是，在同等條件下，如果不使用子進程來執(zhí)行這個函數(shù)，而是直接使用res=f(a)的形式來執(zhí)行，作為臨時變量的b最終依然存在于顯存之中，這是一個非?？膳碌氖虑椤?/p>

4.總結(jié)概要

在使用一些python的GPU模塊，或者寫CUDA時，有時會發(fā)現(xiàn)顯存被無端占用的場景，即時執(zhí)行了cudaFree()或者python的del操作，也無法消除這一塊的顯存占用。最終我們發(fā)現(xiàn)，可以通過額外開啟一個子進程的方法來封裝相關(guān)的操作，通過對進程的存活控制來實現(xiàn)對GPU顯存占用的控制，有可能是一個臨時規(guī)避問題的思路。

到此這篇關(guān)于python中顯存回收問題解決方法的文章就介紹到這了,更多相關(guān)python顯存回收內(nèi)容請搜索腳本之家以前的文章或繼續(xù)瀏覽下面的相關(guān)文章希望大家以后多多支持腳本之家！