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 tf
gpus = 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 的中间输出，引用链未断
验证方法很简单，不用重启进程：
import gc, tracemalloc
tracemalloc.start()
# ... 训练若干步 ...
current, peak = tracemalloc.get_traced_memory()
print(f"Python 分配峰值: {peak / 1024**2:.1f} MB")
gc.collect() # 强制回收，看是否回落如果 peak 持续上涨且 gc.collect() 后不降，就该检查 Python 对象生命周期了，别只盯着 tf.config。
Profiler 能看到什么，又看不到什么
tf.profiler（配合 TensorBoard Profile 标签页）是唯一能定位“哪层 kernel 占显存最多”“哪个 op 分配了巨量临时 buffer”的工具，但它有硬性门槛：
- 必须设好
CUDA_VISIBLE_DEVICES和LD_LIBRARY_PATH包含 CUPTI（/usr/local/cuda/extras/CUPTI/lib64） - Docker 里要加
--privileged=true或至少--cap-add=SYS_ADMIN，否则报CUPTI_ERROR_INSUFFICIENT_PRIVILEGES - 它只抓 profiling 时间段内的快照，不能长期驻留监控；想看趋势还得靠
get_memory_info()定期采样
Profile 页面里的 “Memory Profile” 子页，显示的是 kernel launch 时申请的 device memory，不包括 host memory、Python 对象、或 TF 内部 allocator 的碎片——这些都得靠其他手段交叉验证。
真正难的不是怎么开 profiler，而是当它指出某个 Conv2D op 分配了 2.1GB 临时显存时，你得判断：这是输入 shape 确实太大，还是 pad 方式触发了低效算法，抑或 batch 内图像尺寸不一致导致 padding 补零爆炸。这些细节，API 不会告诉你。
