Question

我想修改现有模型并测试预测。因此，我构建一个图形，对其进行测试，然后构建下一个图形。我在for循环中执行此操作。更详细地讲，在get_new_graph()中，我加载了预先训练的VGG16模型，在其中向网络添加了一层。取决于我选择的test，最后一层的大小会有所不同。

import vgg
slim = tf.contrib.slim

def experiment():
    for test in tests:
        tf.reset_default_graph()
        X, new_pred = get_new_graph(test) # load VGG16 model + add layer
        variables_to_restore = slim.get_variables_to_restore()
        saver = tf.train.Saver(variables_to_restore)
        with tf.Session() as sess:
            saver.restore(sess, './vgg16.ckpt')
            for k in range(100):
                R = sess.run(new_pred, feed_dict={X:images})
                print(R)
            sess.close()

运行此代码时，我可以对imagenet中的1000张图像进行三个测试。然后由于GPU内存已满，我得到了一个内存错误：

W tensorflow/core/common_runtime/bfc_allocator.cc:267]
Allocator (GPU_0_bfc) ran out of memory trying to 
allocate 64.00MiB.  Current allocation summary follows.

如何修改我的代码以使其运行？

Answer 1

正如本期在TensorFlow的github上提到的： http://github.com/tensorflow/tensorflow/issues/17048

似乎可以在不同的进程中创建每个会话，以便在进程终止时释放gpu。

它可能类似于：

from multiprocessing import Pool

def _process(X, new_pred, images):
    with tf.Session() as sess:
        saver.restore(sess, './vgg16.ckpt')
        for k in range(100):
            R = sess.run(new_pred, feed_dict={X:images})
            print(R)
        sess.close()


def experiment():
    for test in tests:
        tf.reset_default_graph()
        X, new_pred = get_new_graph(test) # load VGG16 model + add layer
        variables_to_restore = slim.get_variables_to_restore()
        saver = tf.train.Saver(variables_to_restore)
        with Pool(1) as p:
            return p.apply(_process, (X, new_pred, images,))

Tensorflow：在循环中构建图形时的内存错误

1 个答案: