Tensorflow图像分类 - 耗尽我的记忆

时间:2018-05-14 10:48:09

标签: python-3.x image-processing tensorflow

我正在尝试编写一个能够将某些图片分类的小程序。我在主代码中创建了一个带有图片的列表,并将它们传递给循环中的函数。代码工作得很好,除了它不会释放我的内存,每次迭代程序都会使用更多,直到它完全崩溃。

已经尝试在函数中使用“gc.collect()”来强制它清除内存,但这没有用。检查一个文件后,是否应该自动清除内存,或者我在这里错过了什么?

def classify_pictures(self, files):

    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'

        # Read the image_data
        image_data = tf.gfile.FastGFile(files, 'rb').read()

        # Loads label file, strips off carriage return
        label_lines = [line.rstrip() for line
                   in tf.gfile.GFile("tf_files/retrained_labels.txt")]

        # Unpersists graph from file
        with tf.gfile.FastGFile("tf_files/retrained_graph.pb", 'rb') as f:
            graph_def = tf.GraphDef()
            graph_def.ParseFromString(f.read())
            _ = tf.import_graph_def(graph_def, name='')

        with tf.Session() as sess:
            # Feed the image_data as input to the graph and get first prediction
            softmax_tensor = sess.graph.get_tensor_by_name('final_result:0')

            predictions = sess.run(softmax_tensor, \
                               {'DecodeJpeg/contents:0': image_data})

            # Sort to show labels of first prediction in order of confidence
            top_k = predictions[0].argsort()[-len(predictions[0]):][::-1]

        for each_picture in range(0, 10):

            human_string = label_lines[top_k[0]]

            if human_string == "selfie":
                return ("selfie")

            if "passport" in human_string:
                return("passport")

            if "statement" in human_string:
                return("bill")

1 个答案:

答案 0 :(得分:0)

如果你在循环中调用这个函数,计算图就会在每次迭代时重建,前一个图当然也在那里。这就是使用所有内存的原因。

要解决此问题,只需在循环内进行session.run()调用即可。

通常,在编写Tensorflow代码时,您应该始终尝试将生成图形的代码与执行图形的代码分开。在你的情况下,你在一个函数内部进行,这个函数被多次调用,因此每次都会重建一个新的图形。