不同数据循环使用同一模型tensorflow-gpu计算图后自动释放显存

最新推荐文章于 2025-11-10 10:11:35 发布

原创

最新推荐文章于 2025-11-10 10:11:35 发布 · 2.3k 阅读

3 ·

CC 4.0 BY-SA版权

在TensorFlow中，GPU内存不会在模型执行后自动释放，导致循环使用同一模型时内存持续占用。目前解决方法是在每次运行模型后关闭当前进程，并在新的进程中重新运行计算图，以此达到释放显存的目的。

https://stackoverflow.com/questions/39758094/clearing-tensorflow-gpu-memory-after-model-execution

由于tensorflow中使用完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.