执行模型后清除Tensorflow GPU内存

  • Z时代
  • 2024-01-10
  • 分类：问答

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

加载第一个模型时，它会预先分配整个GPU内存（我希望通过它来处理第一批数据）。但是它不会在完成时卸载内存。当第二模型被加载，同时使用tf.reset_default_graph()和with

tf.Graph().as_default()GPU存储器仍从第一模型完全消耗，并且所述第二模型然后饥饿的存储器。

除了使用Python子进程或多进程来解决该问题（我通过Google搜索找到的唯一解决方案）以外，还有其他方法可以解决此问题吗？

回答：

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

当前，GPUDevice中的分配器属于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），则函数调用后不会释放内存。

以上是 执行模型后清除Tensorflow GPU内存 的全部内容， 来源链接： utcz.com/qa/409257.html

回到顶部