Python programming to view gpu memory usage

Use the package 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,)))

Improvement: Changed to a class, so that the instantiated function can be passed as a parameter

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())