python - 模型执行后清除 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 ) 表明存在以下问题:

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