패키지 pynvml 사용
import pynvml
# ======================GPU========================#
# 查看torch使用gpu情况,需要安装一个包pynvml
# 直接使用pip可以安装
# ======================GPU========================#
def get_gpu_info(use_index=(0,)):
"""
深度学习训练使用,可以回去显卡信息,
使用到的包:pynvml
:param use_index: 使用的GPU的物理编号
:return:
"""
# 计算显存是GB还是MB的函数,方便后续查看数据
def func(number):
# number单位是MB
if number // 1024 > 0: # 如果number对1024取整是大于0的说明单位是GB
return f"{
number / 1024.0:.3f}GB" # 返回值的单位是GB
else:
return f"{
number:.3f}MB"
# 初始化管理工具
pynvml.nvmlInit()
# device = torch.cuda.current_device() # int
gpu_count = pynvml.nvmlDeviceGetCount() # int
information = []
for index in range(gpu_count):
# 不是使用的gpu,就剔除
if index not in use_index:
continue
handle = pynvml.nvmlDeviceGetHandleByIndex(index)
meminfo = pynvml.nvmlDeviceGetMemoryInfo(handle)
total = meminfo.total / 1024 ** 2 # 总的显存大小,单位是MB
used = meminfo.used / 1024 ** 2 # 已用显存大小
free = meminfo.free / 1024 ** 2 # 剩余显存大小
information.append(f"Memory Total:{
func(total)}; Memory Used:{
func(used)}; Memory Free:{
func(free)}")
# 关闭管理工具
pynvml.nvmlShutdown()
return "\n".join(information)
print(get_gpu_info(use_index=(0,)))
개선: 인스턴스화된 함수를 매개변수로 전달할 수 있도록 클래스로 변경
import pynvml
class GetGPUInfo:
"""
# ======================GPU========================#
# 查看torch使用gpu情况,需要安装一个包pynvml
# 直接使用pip可以安装
# ======================GPU========================#
"""
def __init__(self, use_index=(0,)):
self.use_index = use_index
@staticmethod
def get_gpu_info(use_index=(0,)) -> str:
"""
深度学习训练使用,可以回去显卡信息,
使用到的包:pynvml
:param use_index: 使用的GPU的物理编号
:return: 显存使用的信息str
"""
# 计算显存是GB还是MB的函数,方便后续查看数据
def func(number):
# number单位是MB
if number // 1024 > 0: # 如果number对1024取整是大于0的说明单位是GB
return f"{
number / 1024.0:.3f}GB" # 返回值的单位是GB
else:
return f"{
number:.3f}MB"
# 初始化管理工具
pynvml.nvmlInit()
# device = torch.cuda.current_device() # int
gpu_count = pynvml.nvmlDeviceGetCount() # int
information = []
for index in range(gpu_count):
# 不是使用的gpu,就剔除
if index not in use_index:
continue
handle = pynvml.nvmlDeviceGetHandleByIndex(index)
meminfo = pynvml.nvmlDeviceGetMemoryInfo(handle)
total = meminfo.total / 1024 ** 2 # 总的显存大小,单位是MB
used = meminfo.used / 1024 ** 2 # 已用显存大小
free = meminfo.free / 1024 ** 2 # 剩余显存大小
information.append(f"\nMemory Total:{
func(total)}; Memory Used:{
func(used)}; Memory Free:{
func(free)}")
# 关闭管理工具
pynvml.nvmlShutdown()
return "".join(information)
def __call__(self):
return self.get_gpu_info(use_index=self.use_index)
if __name__ == "__main__":
# 使用的显卡的物理编号,用于显示显卡使用信息
use_cuda_index = (0,)
get_gpu_info = GetGPUInfo(use_index=use_cuda_index)
print(get_gpu_info())