解决TensorFlow训练时内存持续增加并占满的问题

周予安

发布时间：2026-10-13| 7人浏览过

原创

本文针对TensorFlow在遍历图片进行风格迁移时内存持续增长的问题，分析了因在循环中重复调用tf计算导致计算图节点累积的根源。通过对比错误与正确代码示例，指出应将所有计算操作定义在模型构建阶段，利用Placeholder统一输入，避免在迭代过程中动态添加新节点，从而有效防止内存溢出。

问题背景与现象描述

本文记录了一次TensorFlow的学习实践，旨在为遇到类似困扰的初学者提供经验参考。

首先回顾错误发生及解决的过程。在进行风格迁移实验时，使用了预加载权重的VGG19网络正向提取中间层结果。由于代码实现不当，在遍历图片提取特征的过程中，内存使用量持续增加，导致仅处理几十张图片后内存即被耗尽。

内存泄漏原因分析

问题的根源在于每张图片正向传播结束后，中间信息未被正确释放。具体而言，由于将正向传播代码与模型定义分离，每次遍历图片时都会重新执行正向传播，这在TensorFlow中新增了大量计算节点（如tf.matmul等），导致内存中遗留了大量过期的计算图信息。

解决方案与代码示例

修正的核心原则是：避免在遍历图片的循环中额外使用TensorFlow的计算资源。

正确的做法是使用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计算节点。

以上关于解决TensorFlow训练时内存持续增加并占满问题的内容到此结束，希望能对大家有所帮助，欢迎支持本站。