标签：显存 程序 gpu 033 print GPU inf 0m

使用之前只需要简单的设置两个变量即可：　

1. Memory_Allocated_GB = 5  # GiB，需要占用的显存大小，单位GiB

2. GPU_ID = 2　　# 需要卡显存的GPU ID

import torch
from torch import cuda
import time
import psutil


def get_gpu_used_information():
    device_count = cuda.device_count()
    used_information = []
    for gpu_id in range(device_count):
        gpu_processes_list = cuda.list_gpu_processes(gpu_id).split('\n')
        for process in gpu_processes_list:
            if process.startswith('process'):
                inf = process.split()
                used_information.append((gpu_id, inf[1], inf[3][:-4], psutil.Process(int(inf[1])).username().split('\\')[1]))

    return used_information


def print_information(infs):
    print('\033[0:32m+----------------------------------------------------------+\033[0m')
    print('\033[0:32m|     GPU       PID       MEMORY-USED       USER-NAME      |\033[0m')
    for inf in infs:
        record = "{:>8d} {:>9d} {:>14d}MiB {}".format(inf[0], int(inf[1]), int(inf[2]), str(inf[3]).rjust(15))
        print('\033[0:32m|' + record + ' ' * 6 + '|\033[0m')
    print('\033[0:32m+----------------------------------------------------------+\033[0m')


if __name__ == '__main__':
    Memory_Allocated_GB = 5  # GiB
    GPU_ID = 2
    ace = None
    while True:
        print('\033[0:33mNow: ' + time.strftime('%Y-%m-%d %H:%M:%S') + '\033[0m')
        used_inf = get_gpu_used_information()
        print_information(used_inf)
        try:
            device = 'cuda:{}'.format(GPU_ID)
            ace = torch.zeros([Memory_Allocated_GB, 1024, 1024, 256], device=device)
        except RuntimeError:
            ace = None
            for sec in range(10)[::-1]:
                print("\r\033[0:31mRace to control GPU: {} {}GiB failed, tra again after {}s...\033[0m".format(
                    GPU_ID, Memory_Allocated_GB,  sec),  flush=True, end='')
                time.sleep(1)
            print()
        else:
            print("\033[1:35mGPU: {}, memory allocated: {} GB\033[0m".format(GPU_ID, Memory_Allocated_GB))
            while True:
                print("\r\033[1:35mGPU: {}, hold-up time: {}\033[0m".format(GPU_ID, time.strftime('%Y-%m-%d %H:%M:%S')), end='',
                      flush=True)
                time.sleep(1)

标签：显存,程序,gpu,033,print,GPU,inf,0m
来源： https://www.cnblogs.com/cpxlll/p/15820619.html

本站声明： 1. iCode9 技术分享网（下文简称本站）提供的所有内容，仅供技术学习、探讨和分享；
2. 关于本站的所有留言、评论、转载及引用，纯属内容发起人的个人观点，与本站观点和立场无关；
3. 关于本站的所有言论和文字，纯属内容发起人的个人观点，与本站观点和立场无关；
4. 本站文章均是网友提供，不完全保证技术分享内容的完整性、准确性、时效性、风险性和版权归属；如您发现该文章侵犯了您的权益，可联系我们第一时间进行删除；
5. 本站为非盈利性的个人网站，所有内容不会用来进行牟利，也不会利用任何形式的广告来间接获益，纯粹是为了广大技术爱好者提供技术内容和技术思想的分享性交流网站。