[파이썬] tensorflow로 GPU 상태(메모리 등) 확인하기

from tensorflow.python.client import device_lib

print(device_lib.list_local_devices())

 

[name: "/device:CPU:0"
device_type: "CPU"
memory_limit: 268435456
locality {
}
incarnation: 12955634870229076918
xla_global_id: -1
, name: "/device:GPU:0"
device_type: "GPU"
memory_limit: 7377227776
locality {
  bus_id: 1
  links {
  }
}
incarnation: 2459137103494862056
physical_device_desc: "device: 0, name: NVIDIA GeForce RTX 3080, pci bus id: 0000:01:00.0, compute capability: 8.6"
xla_global_id: 416903419
]

 

 

import subprocess
import json


def get_gpu_info(nvidia_smi_path='nvidia-smi', no_units=False):
    keys = (
        'index',
        'uuid',
        'name',
        'timestamp',
        'memory.total',
        'memory.free',
        'memory.used',
        'utilization.gpu',
        'utilization.memory'
    )

    nu_opt = '' if not no_units else ',nounits'
    cmd = '%s --query-gpu=%s --format=csv,noheader%s' % (nvidia_smi_path, ','.join(keys), nu_opt)
    output = subprocess.check_output(cmd, shell=True)
    lines = output.decode().split('\n')
    lines = [ line.strip() for line in lines if line.strip() != '' ]

    tmp = [ { k: v for k, v in zip(keys, line.split(', ')) } for line in lines ]
    
    ss = ''
    for key in tmp[0].keys():
        ss += f"{key}:{tmp[0][key]}"
#         if key == 'memory.total' or key == 'memory.free' or key == 'memory.used':
#             ss += " KB\n"
#         else:
#             ss += "\n"
        ss += "\n"
        
    return ss

import pprint
pprint.pprint(get_gpu_info())