执行模型后清除Tensorflow GPU内存
python
tensorflow
18
0

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

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

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

参考资料:
Stack Overflow
收藏
评论
共 3 个回答
高赞 时间 活跃

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

收藏
评论

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

收藏
评论

我使用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))
收藏
评论
新手导航
  • 社区规范
  • 提出问题
  • 进行投票
  • 个人资料
  • 优化问题
  • 回答问题

关于我们

常见问题

内容许可

联系我们

@2020 AskGo
京ICP备20001863号