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产品图添加背景

相关推荐
我是一颗柠檬3 小时前
【MySQL全面教学】MySQL面试高频考点汇总Day15(2026年)
数据库·后端·mysql·面试
星空椰3 小时前
Python 面向对象高级:继承与类定义详解
开发语言·python
橙淮3 小时前
并发编程(六)
java·jvm
凯瑟琳.奥古斯特4 小时前
高阶子查询题目精炼
开发语言·数据库·python·职场和发展·数据库开发
身如柳絮随风扬4 小时前
数据库读写分离:从原理到实战,构建高并发系统
数据库·mysql
风之所往_4 小时前
Python 3.4 新特性全面总结
python
EntyIU4 小时前
JVM内存与GC笔记
java·jvm·笔记
太阳上的雨天4 小时前
任何格式的文件转Markdown
python·ai
提笔了无痕5 小时前
RAG存储策略中.md格式的切片与存储怎么处理
数据库·ai·rag