TensorFlow 1.15下如何记录BERT模型每轮训练后的GPU内存使用情况

TensorFlow 1.15本身没有像TF2那样直接的内存查询API，不过可以通过以下两种方式实现需求：

方法一：调用nvidia-smi命令行工具 #

这是最直接的方式，在Python代码里通过系统命令获取GPU内存数据：

• 导入subprocess模块，在每轮训练结束后执行nvidia-smi并解析输出结果。
• 示例代码：

import subprocess
import re

def get_gpu_memory_usage():
    result = subprocess.check_output(
        ["nvidia-smi", "--query-gpu=memory.used", "--format=csv,nounits,noheader"],
        encoding="utf-8"
    )
    # 解析输出，获取每个GPU的已用内存（单位：MB）
    memory_used = [int(x) for x in result.strip().split("\n")]
    return memory_used

# 在每轮epoch结束后调用
after_epoch_memory = get_gpu_memory_usage()
print(f"Epoch结束后GPU已用内存：{after_epoch_memory} MB")

• 优点：无需依赖TensorFlow内部API，对新手友好，结果准确。

方法二：使用TensorFlow contrib模块的内存统计API #

通过tf.Session调用tf.contrib.memory_stats接口获取内存数据：

• 示例代码：

import tensorflow as tf

def get_tf_gpu_memory(session):
    # 查询当前会话的GPU内存峰值使用量
    memory_stats = session.run(tf.contrib.memory_stats.MaxBytesInUse())
    # 转换为MB（1MB=1024*1024字节）
    memory_used_mb = memory_stats / (1024 * 1024)
    return memory_used_mb

# 假设已创建训练会话
with tf.Session() as sess:
    # 训练循环代码...
    # 每轮epoch结束后调用
    after_epoch_memory = get_tf_gpu_memory(sess)
    print(f"Epoch结束后GPU已用内存：{after_epoch_memory:.2f} MB")

• 注意：tf.contrib.memory_stats属于TensorFlow 1.x的contrib模块，在1.15版本中可正常使用，但不属于核心API范畴。

如果使用的是BERT官方训练代码，只需将上述函数插入到训练循环的epoch结束节点，即可自动记录每轮的GPU内存使用数据。

内容的提问来源于stack exchange，提问作者yama