模型执行后清除 Tensorflow GPU 内存

Clearing Tensorflow GPU memory after model execution

我已经训练了 3 个模型,现在 运行 代码按顺序加载 3 个检查点中的每一个并使用它们运行预测。我正在使用 GPU。

当加载第一个模型时,它会预先分配整个 GPU 内存(我想要处理第一批数据)。但它不会在完成后卸载内存。当加载第二个模型时,同时使用 tf.reset_default_graph()with tf.Graph().as_default(),第一个模型的 GPU 内存仍然被完全消耗,然后第二个模型内存不足。

除了使用 Python 子进程或多进程来解决这个问题(我通过 google 搜索找到的唯一解决方案)之外,还有其他解决方法吗?

一旦不再需要张量(在 .运行 调用终止之前),张量分配的 GPU 内存就会被释放(返回到 TensorFlow 内存池中)。当变量容器被销毁时,为变量分配的 GPU 内存将被释放。在 DirectSession 的情况下(即 sess=tf.Session("")),它是在会话关闭或显式重置时(在 62c159ff 中添加)

2016 年 6 月的 git 问题 (https://github.com/tensorflow/tensorflow/issues/1727) 表示存在以下问题:

currently the Allocator in the GPUDevice belongs to the ProcessState, which is essentially a global singleton. The first session using GPU initializes it, and frees itself when the process shuts down.

因此唯一的解决方法是使用进程并在计算后关闭它们。

示例代码:

import tensorflow as tf
import multiprocessing
import numpy as np

def run_tensorflow():

    n_input = 10000
    n_classes = 1000

    # Create model
    def multilayer_perceptron(x, weight):
        # Hidden layer with RELU activation
        layer_1 = tf.matmul(x, weight)
        return layer_1

    # Store layers weight & bias
    weights = tf.Variable(tf.random_normal([n_input, n_classes]))


    x = tf.placeholder("float", [None, n_input])
    y = tf.placeholder("float", [None, n_classes])
    pred = multilayer_perceptron(x, weights)

    cost = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=pred, labels=y))
    optimizer = tf.train.AdamOptimizer(learning_rate=0.001).minimize(cost)

    init = tf.global_variables_initializer()

    with tf.Session() as sess:
        sess.run(init)

        for i in range(100):
            batch_x = np.random.rand(10, 10000)
            batch_y = np.random.rand(10, 1000)
            sess.run([optimizer, cost], feed_dict={x: batch_x, y: batch_y})

    print "finished doing stuff with tensorflow!"


if __name__ == "__main__":

    # option 1: execute code with extra process
    p = multiprocessing.Process(target=run_tensorflow)
    p.start()
    p.join()

    # wait until user presses enter key
    raw_input()

    # option 2: just execute the function
    run_tensorflow()

    # wait until user presses enter key
    raw_input()

因此,如果您在创建的进程中调用函数 run_tensorflow() 并关闭该进程(选项 1),内存将被释放。如果你只是 运行 run_tensorflow() (选项 2)函数调用后内存不会被释放。

现在似乎有两种方法可以解决迭代训练模型或者如果你使用未来的多进程池来为模型训练服务,如果未来完成,池中的进程不会被杀死。您可以在训练过程中使用两种方法来释放GPU内存,同时您希望保留主进程。

  1. 调用子进程以 运行 模型训练。当一个阶段训练完成时,子进程将退出并释放内存。很容易得到return值。
  2. 调用multiprocessing.Process(p)到运行模型训练(p.start),p.join将指示进程退出并释放内存。

这是一个使用 multiprocess.Process 的辅助函数,它可以为 运行 您的 python 编写的函数和返回值打开一个新进程,而不是使用子进程,

# open a new process to run function
def process_run(func, *args):
    def wrapper_func(queue, *args):
        try:
            logger.info('run with process id: {}'.format(os.getpid()))
            result = func(*args)
            error = None
        except Exception:
            result = None
            ex_type, ex_value, tb = sys.exc_info()
            error = ex_type, ex_value,''.join(traceback.format_tb(tb))
        queue.put((result, error))

    def process(*args):
        queue = Queue()
        p = Process(target = wrapper_func, args = [queue] + list(args))
        p.start()
        result, error = queue.get()
        p.join()
        return result, error  

    result, error = process(*args)
    return result, error

我使用numba释放GPU。使用TensorFlow,我找不到有效的方法。

import tensorflow as tf
from numba import cuda

a = tf.constant([1.0,2.0,3.0],shape=[3],name='a')
b = tf.constant([1.0,2.0,3.0],shape=[3],name='b')
with tf.device('/gpu:1'):
    c = a+b

TF_CONFIG = tf.ConfigProto(
gpu_options=tf.GPUOptions(per_process_gpu_memory_fraction=0.1),
  allow_soft_placement=True)

sess = tf.Session(config=TF_CONFIG)
sess.run(tf.global_variables_initializer())
i=1
while(i<1000):
        i=i+1
        print(sess.run(c))

sess.close() # if don't use numba,the gpu can't be released
cuda.select_device(1)
cuda.close()
with tf.device('/gpu:1'):
    c = a+b

TF_CONFIG = tf.ConfigProto(
gpu_options=tf.GPUOptions(per_process_gpu_memory_fraction=0.5),
  allow_soft_placement=True)

sess = tf.Session(config=TF_CONFIG)

sess.run(tf.global_variables_initializer())
while(1):
        print(sess.run(c))

你可以使用numba库来释放所有的gpu内存

pip install numba 
from numba import cuda 
device = cuda.get_current_device()
device.reset()

这将释放所有内存

我正在找出 Jupyter Notebook 中哪个选项更好。即使深度学习应用程序完成,Jupyter Notebook 也会永久占用 GPU 内存。它通常会导致令人头疼的 GPU Fan ERROR。在这种情况下,我必须定期重置 nvidia_uvm 并重新启动 linux 系统。我得出以下两个选项可以消除 GPU Fan Error 的麻烦,但想知道哪个更好。

环境:

  • CUDA 11.0
  • cuDNN 8.0.1
  • 张量流 2.2
  • 凯拉斯 2.4.3
  • Jupyter 笔记本 6.0.3
  • Miniconda 4.8.3
  • Ubuntu 18.04 LTS

第一个选项

将以下代码放在单元格的末尾。内核在应用程序运行时完成后立即结束。但它并不优雅。 Juputer 将弹出一条关于死机内核的消息。

import os
 
pid = os.getpid()
!kill -9 $pid

部分选项

下面的代码也可以用Jupyter Notebook结束内核。我不知道 numba 是否安全。 Nvidia 更喜欢“0”GPU,这是个人开发人员(不是服务器人员)使用最多的 GPU。然而,Neil G 和 mradul dubey 都得到了回应:这让 GPU 处于糟糕的状态。

from numba import cuda

cuda.select_device(0)
cuda.close()

看来第二种方案比较优雅。有人可以确认哪个是最佳选择吗?

备注:

Anaconda环境下直接执行"$pythonabc.py自动释放GPU内存不是这个问题。然而,我有时需要使用Jyputer Notebook 来处理.ipynb 应用程序。

当我训练了 120 个模型后出现此错误时,我已经在 for 循环中针对不同的参数训练了我的模型。之后如果我不杀掉内核,我什至无法训练一个简单的模型。 我能够通过在构建模型之前添加以下行来解决我的问题:

tf.keras.backend.clear_session()

(参见 https://www.tensorflow.org/api_docs/python/tf/keras/backend/clear_session