TensorFlow如何监控内存使用情况_结合tf.summary记录关键指标信息

tf.config.experimental.get_memory_info() 是 TensorFlow 2.15+ 最轻量的实时显存监控方式，仅返回已分配显存（current），需先启用 set_memory_growth，不支持 ROCm 且无法反映 Python 层内存泄漏。如何用 tf.config.experimental.get_memory_info() 实时查显存这是 TensorFlow 2.15+ 最轻量、最直接的 GPU 显存监控方式，不依赖 TensorBoard，也不需要改模型结构。但它只返回"已分配"显存（current），不是"GPU总占用"，更不等于 nvidia-smi 看到的 used memory------因为 CUDA 驱动层还管着底层缓存和未释放的临时张量。必须先启用内存增长，否则 get_memory_info() 可能返回 0 或异常值：import tensorflow as tfgpus = tf.config.experimental.list_physical_devices('GPU')if gpus: for gpu in gpus: tf.config.experimental.set_memory_growth(gpu, True) # 关键！ mem_info = tf.config.experimental.get_memory_info('GPU:0') print(f"当前分配显存: {mem_info['current'] / 1024**3:.2f} GB")只支持 'GPU:0' 这类逻辑设备名，不能传 gpu.name若报 NotImplementedError，说明当前后端（如 ROCm）不支持该接口，得切回 CUDA 环境它不反映 Python 层对象引用导致的内存泄漏，比如 tf.data.Dataset 中闭包捕获了大数组为什么 tf.summary.scalar() 记不了真实显存峰值tf.summary.scalar() 本身不能直接记录 get_memory_info() 的结果，因为 summary ops 必须在图执行上下文中（@tf.function 内或 eager 模式下明确调用 tf.summary.record_if(True)），而 get_memory_info() 是纯 Python 调用，无法被自动追踪进计算图。所以你不能写 tf.summary.scalar('gpu_mem', mem_info['current']) 并期望它自动出现在 TensorBoard 里------它会静默失败，或者报 Outside of tf.function 错误。正确做法是：在训练循环中手动取值 + 手动写 summary writer必须用 tf.summary.create_file_writer() 和 as_default() 上下文记得每步都调用 writer.flush()，否则 TensorBoard 刷新延迟高、甚至漏数据train_summary_writer = tf.summary.create_file_writer('./logs/train')# 在 @tf.function 外的训练 loop 中：for step, (x, y) in enumerate(dataset): train_step(x, y) if step % 100 == 0: mem = tf.config.experimental.get_memory_info('GPU:0')['current'] with train_summary_writer.as_default(): tf.summary.scalar('gpu_mem_allocated_GB', mem / 1024**3, step=step) train_summary_writer.flush()显存"缓慢上涨"却没报 OOM？大概率是 Python 层泄漏训练跑几小时后 get_memory_info() 显示从 3GB 涨到 7GB，但模型 batch size 和结构完全没变------这往往不是 TensorFlow 张量泄漏，而是 Python 对象堆积，比如：自定义 tf.data.Dataset.from_generator() 中，生成器函数里不断追加 list / dict，且没清空数据增强用了 OpenCV/PIL 加载图像但没 del img，又没触发 gc回调（Callback）里缓存了每个 batch 的中间输出，引用链未断验证方法很简单，不用重启进程： Mokker AI AI产品图添加背景