[解決済み] モデル実行後にTensorflow GPUのメモリをクリアする

質問

3つのモデルをトレーニングし、現在、3つのチェックポイントを順番にロードし、それらを使って予測を実行するコードを実行しているところです。GPUを使用しています。

最初のモデルがロードされると、GPUメモリ全体が事前に割り当てられます(これは最初のデータバッチを処理するために必要です)。しかし、それが終わっても、メモリはアンロードされません。2番目のモデルがロードされるとき、両方の tf.reset_default_graph() と with tf.Graph().as_default() の場合、GPU メモリはまだ最初のモデルから完全に消費されており、2 番目のモデルはメモリが不足します。

Pythonのサブプロセスやマルチプロセッシングを使用して問題を回避する以外に、これを解決する方法はありますか（グーグル検索で見つけた唯一の解決策です）。

解決方法は？

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) 関数呼び出しの後、メモリは解放されません。