Tensorflow 程序内存泄漏+显存分配不足数据预处理

最新推荐文章于 2025-06-13 09:59:48 发布

原创

最新推荐文章于 2025-06-13 09:59:48 发布 · 2.9k 阅读

2 ·

CC 4.0 BY-SA版权

文章标签：

#图像batch #内存泄漏 #显存使用增加 #tensorflow

本文旨在解决Tensorflow程序中的内存泄漏和显存分配不足问题。错误步骤是由于在循环中直接处理tensor导致的。提出了两种修改方法：1) 图片预处理阶段避免使用tensor，改用numpy；2) 利用tf.train.batch等内置函数优化。通过调试测试函数监控显存使用和变量引用次数来验证效果。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

1. 目的

在循环中，从硬盘中读取每个batch 的图片。

2. 错误步骤及原因

图片截取从here。

3. 修改方法

方法一：可以将图片数据预处理部分不使用 tensor ,用的python 常见的numpy 等类型，这样就不会产生新的tensor 结点。
方法二：使用 tf.train.batch 等函数。

4. 调试测试函数

查看显卡显存使用量

import pynvml
pynvml.nvmlInit()

handle = pynvml.nvmlDeviceGetHandleByIndex(0)     #   第0块显卡
meminfo = pynvml.nvmlDeviceGetMemoryInfo(handle)
print(meminfo.used)                               # 查看显卡显存使用量

查看内存变量类型引用次数（当前次数，增加次数）

import gc
import objgraph

gc.collect()       # 变量回收
objgraph.show_growth()                     #显示当前变量类型引用次数 及增加引用次数
objgraph.show_most_common_types(limit=50)  #显示引用次数最多的50个变量类型