\TOC@[ ]
显存状态的获取对于深度学习任务至关重要，特别是在使用 GPU 进行模型训练和推理时。了解当前显存的使用情况可以帮助我们优化算法、调整超参数，以及有效地管理计算资源。在本文中，我们将讨论多种获取当前显存状态的方法，并探讨它们的准确性和原理。
获取显存开销的方式：
pytorch的cuda相关apipytroch分析工具nvidia-smipycuda
nvidia-smi
首先来看最简单的，官方工具：nvidia-smi。最简单，并且只要安装了驱动，就一定可用。显示内容也很丰富，包括占用进程、显存、风扇等。
nvidia-smi 是基于 NVIDIA 驱动程序提供的接口实现的。这些接口允许用户查询在 GPU 内部的寄存器中各种指标数据，如 GPU 温度、GPU 使用率、显存使用情况等。
可以使用kill -9 PID号来杀掉未关闭的程序
可以在linux使用：
watch -n 1 nvidia-smi
或者在wondws使用
nvidia-smi -l 5
来持续获取显存
Pytorch
在 PyTorch 中，可以使用 torch.cuda.memory_allocated() 和 torch.cuda.memory_reserved() 函数来查询当前 GPU 显存的使用情况。这两个函数提供了获取已分配显存和已保留显存的功能，但并不直接获取当前实际的显存使用情况。
两个函数可能导致与实际显存不一致。例如下面的情况：
import torch
# 检查CUDA是否可用
if torch.cuda.is_available():
# 获取CUDA设备数量
device_count = torch.cuda.device_count()
print("CUDA可用，共有 {} 个CUDA设备可用:".format(device_count))
for i in range(device_count):
device = torch.device("cuda:{}".format(i))
print("CUDA 设备 {}: {}".format(i, torch.cuda.get_device_name(i)))
# 获取当前设备的显存使用情况
total_memory = torch.cuda.get_device_properties(device).total_memory
allocated_memory = torch.cuda.memory_allocated(device) # 已分配的显存
reserved_memory = torch.cuda.memory_reserved(device) # 已保留的显存
free_memory = total_memory - allocated_memory - reserved_memory # 剩余可用显存
print(" - 总显存: {:.2f} GB".format(total_memory / (1024 ** 3)))
print(" - 已分配的显存: {:.2f} GB".format(allocated_memory / (1024 ** 3)))
print(" - 已保留的显存: {:.2f} GB".format(reserved_memory / (1024 ** 3)))
print(" - 剩余可用显存: {:.2f} GB".format(free_memory / (1024 ** 3)))
else:
print("CUDA 不可用")
torch.cuda.memory_allocated()： 函数返回已经由 PyTorch 分配但尚未释放的显存字节数。它只能反映出 PyTorch 所管理的显存使用情况，即已被张量、模型参数等数据占用的显存量。
torch.cuda.memory_reserved()： 函数返回由 PyTorch 预先保留但尚未实际分配的显存字节数。PyTorch 在启动时会预留一定数量的显存作为缓冲区，以便在需要时快速分配。这部分显存并不会直接被占用，因此不计入已分配显存。
需要注意的是，如果你使用不同的python环境，例如Docker、Conda，这样的命令无法获取其他python进程的显存开销。是不准确的。
PyCUDA
PyCUDA 是 CUDA 的 Python 接口，它通过调用 CUDA 库中的函数来实现显存的管理和操作。cuda.mem_get_info() 函数实际上是调用了 CUDA 库中的相应函数来获取显存使用情况的信息。
在 PyCUDA 中，要查看 GPU 显存的使用情况，可以使用 cuda.mem_get_info() 函数。这个函数返回一个元组，其中包含两个元素：可用显存的字节数和已分配显存的字节数。
cuda.mem_get_info() 函数： 这个函数用于查询当前 GPU 设备的显存使用情况。它会返回一个包含两个元素的元组，分别表示当前可用的显存字节数和已分配的显存字节数。
用户调用 cuda.mem_get_info() 函数时，PyCUDA 会通过调用 CUDA 库中的相应函数来访问 GPU 设备的监控数据。CUDA 库会通过 PCI 总线访问 GPU 设备，并读取相应的寄存器中的数据。
import pycuda.driver as cuda
# 初始化 PyCUDA
cuda.init()
# 获取 GPU 数量
device_count = cuda.Device.count()
print("GPU 可用，共有 {} 个 GPU 设备可用:".format(device_count))
for i in range(device_count):
# 获取 GPU 设备
device = cuda.Device(i)
print("GPU 设备 {}: {}".format(i, device.name()))
# 创建上下文
context = device.make_context()
# 查询显存使用情况
total_memory = device.total_memory()
free_memory = cuda.mem_get_info()[0]
allocated_memory = total_memory - free_memory
# 打印显存使用情况
print(" - 总显存: {:.2f} GB".format(total_memory / (1024 ** 3)))
print(" - 已分配的显存: {:.2f} GB".format(allocated_memory / (1024 ** 3)))
print(" - 剩余可用显存: {:.2f} GB".format(free_memory / (1024 ** 3)))
# 释放上下文
context.pop()
结果是准确的
gpustat
这个命令会显示当前系统中所有 GPU 设备的简要信息，包括 GPU 的编号、型号、利用率、显存使用情况等。如果你没有安装 gpustat，可以使用以下命令进行安装：
讲一下显存的种类
参考评论区
在PyTorch中，显存管理是一个复杂但高效的系统，它通过分层的方式管理GPU显存资源。以下是关键概念的详细说明及其关系：
1. 核心概念
reserved_memory（保留显存）
PyTorch通过内存分配器（如CUDA的caching_allocator）预先从GPU保留的显存总量。这部分包含：
memory_allocated：当前已分配给张量/变量的显存（活跃内存）。缓存区（Cache）：未被变量占用但由分配器保留的内存池，用于快速响应未来分配请求，避免频繁向GPU申请/释放显存。
PyTorch context（上下文开销）
PyTorch运行时自身的开销，包括：
CUDA上下文初始化（驱动和硬件通信的基础设施）。内核函数、CUDA库占用的显存（如cuBLAS、cuDNN的静态资源）。内部数据结构（如执行引擎、流管理）的显存占用。
2. 显存层级关系
总显存占用可表示为：
total_memory = reserved_memory + PyTorch_context
reserved_memory
由PyTorch主动控制，通过torch.cuda.memory_reserved()查询。动态调整：当memory_allocated增加时，若缓存不足，会从GPU申请更多显存到reserved_memory。
PyTorch context
在首次CUDA操作时初始化，与框架深度绑定。通常固定（约100MB~1GB，取决于GPU架构和PyTorch版本），不随模型变化。
3. 监控显存的API
torch.cuda.memory_allocated()：当前张量占用的显存（memory_allocated）。torch.cuda.memory_reserved()：PyTorch保留的总显存（reserved_memory）。torch.cuda.max_memory_allocated()：进程生命周期内memory_allocated的峰值。torch.cuda.max_memory_reserved()：reserved_memory的峰值。
4. 示例与场景分析
import torch
# 初始化CUDA，产生PyTorch context
x = torch.randn(1000, 1000, device='cuda')
print(torch.cuda.memory_allocated()) # 变量显存（约4MB）
print(torch.cuda.memory_reserved()) # 保留显存（通常更大，如1GB）
print(torch.cuda.max_memory_allocated()) # 峰值分配量
典型场景：
训练初期：reserved_memory会逐步增长，以缓存足够空间供后续使用。显存碎片化：频繁分配/释放可能导致reserved_memory远大于memory_allocated，此时可调用torch.cuda.empty_cache()释放未使用的缓存。
5. 优化建议
控制缓存：通过torch.cuda.set_per_process_memory_fraction()限制reserved_memory上限。减少上下文开销：避免不必要的CUDA操作（如频繁切换设备）。及时清缓存：在验证/测试阶段手动调用empty_cache()。