GPU如何统计显存占用

GPU统计显存占用有两种方法：

方法1： 使用

watch -n 1 nvidia-smi 

这个会2秒刷新一次，可以在一个终端运行程序另外一个终端执行这个然后看具体的显存变化【很费劲，而且可能一闪而过】

方法2：

要统计程序的最大显存占用，你可以使用Python的nvidia-smi库来监测程序的显存使用情况并记录最大值。以下是一个示例代码，用于实现这个目标：

首先，确保你已经安装了NVIDIA的nvidia-smi工具和相关驱动，以及Python的nvidia-ml-py库，这个库用于与NVIDIA Management Library (NVML) 进行交互。

可以使用以下命令来安装 nvidia-ml-py：

pip install nvidia-ml-py

然后，你可以使用以下Python代码来统计程序的最大显存占用：

import time
import pynvml
 
def get_max_gpu_memory_usage():
    pynvml.nvmlInit()
    
    try:
        gpu_count = pynvml.nvmlDeviceGetCount()
        max_memory_usage = [0] * gpu_count
 
        while True:
            for i in range(gpu_count):
                handle = pynvml.nvmlDeviceGetHandleByIndex(i)
                memory_info = pynvml.nvmlDeviceGetMemoryInfo(handle)
                memory_used = memory_info.used
 
                if memory_used > max_memory_usage[i]:
                    max_memory_usage[i] = memory_used
            
            # Delay to control the monitoring interval
            time.sleep(1)
 
    except KeyboardInterrupt:
        return max_memory_usage
    finally:
        pynvml.nvmlShutdown()
 
if __name__ == "__main__":
    max_memory_usage = get_max_gpu_memory_usage()
    
    for i, max_memory in enumerate(max_memory_usage):
        print(f"Maximum GPU {i} memory usage: {max_memory / 1024 / 1024} MB")

在当前终端运行你的python 程序，在另一个终端运行(也可以本终端后台运行这个）这个GPU显存统计脚本【这个代码会初始化NVML库，然后在一个无限循环中检查每块GPU的显存占用，并记录每块GPU的最大显存占用。你可以通过按Ctrl+C来停止程序的执行，并最终打印出每块GPU的最大显存占用。】

当你的python 程序运行结束后，Ctrl+C来停止本脚本，就会出现各个device的打印结果