Получение памяти графического процессора, используемого текущим скриптом - PullRequest
Новая
       20

Получение памяти графического процессора, используемого текущим скриптом

0 голосов
/ 02 марта 2020

Я внес некоторые изменения в решение, опубликованное ранее sam в этой теме: Найти графический процессор с достаточным объемом памяти 

import subprocess

def get_current_gpu_usage():
    """
    Get the current gpu usage of this script.
    Requires nvidia-smi and subprocess module.

    Returns memory used (MiB)

    """
    result_PID = subprocess.check_output(
        [
            'nvidia-smi', '--query-compute-apps=pid',
            '--format=csv,nounits,noheader'
        ], encoding='utf-8')
    gpu_PIDS = [int(x) for x in result_PID.strip().split('\n')]

    result = subprocess.check_output(
        [
            'nvidia-smi', '--query-compute-apps=used_memory',
            '--format=csv,nounits,noheader'
        ], encoding='utf-8')
    # Convert lines into a list
    gpu_memory = [int(x) for x in result.strip().split('\n')]

    # Get current process ID
    PID = os.getpid()

    # Find index of current process within gpu_PIDS
    try:
        index = gpu_PIDS.index(PID)
    except ValueError:
        return 0

    return gpu_memory[index]

Я хотел бы знать, будет ли какие-либо возможные осложнения при использовании с другим импортом, таким как os? (Я не очень хорошо знаю, как работают subprocess и os). Если нет, я надеюсь, что этот фрагмент кода поможет кому-то с управлением памятью на GPU! (Я использовал это только в среде linux)

PS Может кто-нибудь пролить свет на разницу между MiB и MB? (в nvidia-smi)

...