坑 之 tensorflow使用sess.run处理图片时越来越慢,占用内存越来越大的问题

最近在项目中需要使用tensorflow来对图片进行离线预处理,但是使用sess.run时代码越来越慢,最后几分钟才能处理一张图片,这可怎么能行,我的数据集有数万张图片......

原始代码

def process(img, .. , .. ,):
    
    .......

    return img

def main():
    .......
    with tf.Session() as sess:
        for i in range(len(dataset)):
            img_op = process(...) #构图动作
            img = sess.run(img_op)
    .......

if __name__ == "__main__":
    main()

代码在执行的过程中越来越慢,并且占用的内存空间越来越大,这就说明代码本身是存在问题的。
我们知道,tensorflow在执行之前是要先构图,然后再向图中feed数据,正常来说,这两个过程应该是分开进行的,但是由于我们在编写代码时的不规范,就容易在feed数据的过程在向计算图中添加新的计算节点,你比如,我这里的img_op = process(.....)本身就是添加计算节点的动作,又叫做构图动作,这些动作是非常危险的,不仅会影响运行速度,还会扩大计算图的规模,占用过多的内存,有可能会导致程序崩溃。所以导致我代码运行缓慢的原因是找到了,那么该怎么解决呢?
有两种方案:1.将构图动作移动到循环外。2.再循环里不断变换当前的计算图。
我们来分析以下两种方法分别适合什么样的场景:
1.将构图动作移动到循环外。代码如下:

def process(, .. , .. ,):
    img = tf.placeholder(tf.int64, shape=[,,,], name='input')
    
    .......

    return img

def main():
    .......
    img_op = process(...) #构图动作
    with tf.Session() as sess:
        for i in range(len(dataset)):
            img_np = np.array(..)
            img = sess.run(img_op, feed_dict = {input:0: img_np})
    .......

if __name__ == "__main__":
    main()

使用标准的构图方式建立一个计算图,使用feed_dict喂入数据,那么聪明的你可能已经发现了一个问题,就是,如果我的图片size大小不一怎么办,没办法使用tf.placeholder来占位。没错,这种办法由于是静态计算图的原因,只适用于固定输入尺寸的处理任务,想要解决刚刚的问题只能使用第二种办法。
2.再循环里不断变换当前的计算图。代码如下:

def process(img, .. , .. ,):
    
    .......

    return img

def main():
    .......
    for i in range(len(dataset)):
        tf.reset_default_graph() #########重置计算图
        with tf.Session() as sess:
            img_op = process(...) #构图动作,创建计算图
            img = sess.run(img_op)
            tf.get_default_graph().finalize() #关闭当前计算图
    .......

if __name__ == "__main__":
    main()

这种办法,我们发现,我们每次送入图中的数据都可以是尺寸不一的,因为在处理每一个数据时都会创建新的计算图,并且在处理结束后关闭当前计算图。通过这样的方式,我们就实现了代码维护一个动态的计算图,就避免了一开始计算图越来越大所导致的种种问题。

两种办法所使用的场景不相同,重点是大家可以理解二者的区别究竟在哪里,并理解tensorflow的图计算机制,来选择适合自己工程的方式

猜你喜欢

转载自blog.csdn.net/qq_41368074/article/details/111630636