tensorflow训练时内存持续增加并占满

妖狐艹你老母 2021-11-23 14:46 222阅读 0赞

记录一次小白的tensorflow学习过程，也为有同样困扰的小白留下点经验。

先说我出错和解决的过程。在做风格迁移实验时，使用预加载权重的VGG19网络正向提取中间层结果，结果因为代码不当，在遍历图片提取时内存持续增长，导致提取几十个图片的特征内存就满了。

原因是在对每一张图片正向传播结束后，都会在留下中间信息。具体地说是在我将正向传播的代码与模型的代码分离了，在每次遍历图片时都会正向传播，在tensorflow中新增加了很多的计算节点（如tf.matmul等等），导致内存中遗留了大量的过期信息。

纠正的做法就是一个前提：避免在循环训练图片时额外使用tf计算资源。

使用placeholder作为输入数据的入口，在模型中定义需要使用的函数，包括正向传播。不要在遍历图片时额外使用tf计算。

遇到这种问题一定要回头检查代码，尤其是在别人写的代码基础上改时。 多学习公开的源码。

错误示例：

def build_model(model_path):
        model_input = tf.placeholder('float32', [1, IMAGE_HEIGHT, IMAGE_WIDTH, COLOR_CHANNELS])
        vec1 = ...
        ...... 
        return model_input,vec1
        
    def get_style_represent(vec):
        # 一些tf计算操作
        return new_vec
    
    with tf.Session() as sess:
        sess.run(tf.global_variables_initializer())
        img_input,vec1 = build_model(VGG19_MODEL)       # 加载模型
        for cur_img_path in imgs_path_list:     # 遍历图片
            cur_img = load_image(cur_img_path)
            vec1_out = sess.run(vec1, feed_dict = {img_input:cur_img})      # 正向传播输出模型中的vec1
            # 对vec1进行一些处理，此处在遍历图片时额外使用了tensorflow的计算节点，导致在内存中遗留信息
            new_vec = get_style_represent(vec1_out)

正确实例：

def build_model(model_path):
        model_input = tf.placeholder('float32', [1, IMAGE_HEIGHT, IMAGE_WIDTH, COLOR_CHANNELS])
        vec1 = ...
        ...... 
        new_vec = ...       # 将get_style_represent计算操作定义在模型中
        return model_input,vec1,new_vec
    
    with tf.Session() as sess:
        sess.run(tf.global_variables_initializer())
        img_input,vec1,new_vec = build_model(VGG19_MODEL)
        for cur_img_path in imgs_path_list:
            cur_img = load_image(cur_img_path)
            # 一次正向传播将处理后的vec1也得到了，避免在每次图片正向传播时留下额外信息
            vec1_out,new_vec_out = sess.run([vec1,new_vec], feed_dict = {img_input:cur_img})

所以，如果你也遇到了同样的问题，不妨看一下你是不是在迭代过程中额外使用了新的tf计算节点吧。  
\---------------------   
作者：张小白027   
来源：CSDN   
原文：https://blog.csdn.net/qq\_18714153/article/details/80682387   
版权声明：本文为博主原创文章，转载请附上博文链接！