模型执行后清除Tensorflow GPU内存

Question

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

加载第一个模型时,它会预先分配整个GPU内存(我希望通过第一批数据处理).但是当它完成时它不会卸载内存.当加载第二个模型时,使用两个tf.reset_default_graph()并且with tf.Graph().as_default()GPU内存仍然完全从第一个模型消耗,然后第二个模型缺乏内存.

有没有办法解决这个问题,除了使用Python子进程或多处理来解决问题(我通过谷歌搜索找到的唯一解决方案)？

Answer 1

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

目前GPUDevice中的Allocator属于ProcessState,它本质上是一个全局单例.使用GPU的第一个会话初始化它,并在进程关闭时释放自己.

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

示例代码:

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),则在函数调用后不释放内存.

Answer 2

您可以使用 numba 库释放所有 gpu 内存

pip install numba 

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

这将释放所有内存

Answer 3

我用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))

Answer 4

我刚才用垃圾收集器解决了 OOM 错误。

import gc
gc.collect()

model.evaluate(x1, y1)
gc.collect()

model.evaluate(x2, y2)
gc.collect()

ETC。

根据Yaroslav Bulatov在他们的回答中所说的（当对象被销毁时 tf 会释放 GPU 内存），我猜测可能只是垃圾收集器尚未运行。强迫它收集让我解放了，所以这可能是一个好方法。

Answer 5

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

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

这是一个使用 multiprocess.Process 的辅助函数，它可以打开一个新进程来运行 python 编写的函数并返回值，而不是使用 Subprocess，

# 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