Skip to content

memory

util.memory

CudaMemoryTracker 

CudaMemoryTracker(device: Union[device, str])

특정 구간에서 추가적으로 사용된 최대 GPU 메모리를 측정합니다. 예를 들어 모델을 먼저 GPU에 올린 뒤 __enter__에 진입하면 모델이 점유한 메모리는 제외됩니다. 각 연산의 메모리를 정확하게 트래킹하려면 torch.profiler를 사용하세요.

Example
  1. Context Manager로 사용하는 경우 with CudaMemoryTracker() as tracker: results = model.inference(batch)
  2. class로 사용하는 경우 tracker = CudaMemoryTracker() tracker.start() results = model.inference(batch) tracker.end() ...

print(tracker.get_results())

Source code in SaigeToolkit/util/memory.py 
 95
 96
 97
 98
 99
100
101
102
def __init__(self, device: Union[torch.device, str]):
    self.device = torch.device(device)
    self.enabled = all(
        [self.device.type == "cuda", platform.system() == "Linux", torch.cuda.is_available()]
    )
    self.cuda_context_memory = 0
    self.max_memory_allocated = 0
    self.max_memory_reserved = 0

get_process_cpu_memory 

get_process_cpu_memory(pid: int = None, include_virtual: bool = False) -> int

프로세스가 사용하는 RAM 크기를 바이트 단위로 읽어옵니다.

Parameters:

  • pid (int, default: None ) –

    If None, use current PID.

  • include_virtual (bool, default: False ) –

    to include virtual memory. Defaults to False.

Returns:

  • int ( int ) –

    CPU memory in bytes

Note

Use --pid=host flag with Docker.

Source code in SaigeToolkit/util/memory.py 
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
def get_process_cpu_memory(pid: int = None, include_virtual: bool = False) -> int:
    """프로세스가 사용하는 RAM 크기를 바이트 단위로 읽어옵니다.

    Args:
        pid (int, optional):  If None, use current PID.
        include_virtual (bool, optional): to include virtual memory. Defaults to False.

    Returns:
        int: CPU memory in bytes

    Note:
        Use --pid=host flag with Docker.
    """
    if pid is None:
        pid = os.getpid()
    memory_info = psutil.Process(pid).memory_info()
    if include_virtual:
        return memory_info.vms
    else:
        return memory_info.rss

get_process_gpu_memory 

get_process_gpu_memory(pid: int = None) -> int

프로세스가 사용하는 GPU메모리를 NVML 모듈로부터 바이트 단위로 읽어옵니다.

Parameters:

  • pid (int, default: None ) –

    If None, use current PID.

Returns:

  • int ( int ) –

    GPU memory in bytes

Note

Use --pid=host flag with Docker.

Source code in SaigeToolkit/util/memory.py 
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
def get_process_gpu_memory(pid: int = None) -> int:
    """프로세스가 사용하는 GPU메모리를 NVML 모듈로부터 바이트 단위로 읽어옵니다.

    Args:
        pid (int, optional): If None, use current PID.

    Returns:
        int: GPU memory in bytes

    Note:
        Use --pid=host flag with Docker.
    """
    if pid is None:
        pid = os.getpid()
    pynvml.nvmlInit()
    for device_id in range(pynvml.nvmlDeviceGetCount()):
        handle = pynvml.nvmlDeviceGetHandleByIndex(device_id)
        for process in pynvml.nvmlDeviceGetComputeRunningProcesses(handle):
            if process.pid == pid:
                if process.usedGpuMemory is None:
                    return 0
                else:
                    return process.usedGpuMemory
    return 0

get_memory_stats 

get_memory_stats() -> Dict[str, int]

현재 프로세스의 메모리 정보를 MiB 단위로 읽어옵니다. torch_max_allocated gpu mem의 경우 해당 함수 호출 직전까지의 최댓값이며, 해당 함수 호출 시 리셋됩니다.

Returns:

  • Dict[str, int] –

    Dict[str, int]: description

Source code in SaigeToolkit/util/memory.py 
60
61
62
63
64
65
66
67
68
69
70
71
72
73
def get_memory_stats() -> Dict[str, int]:
    """현재 프로세스의 메모리 정보를 MiB 단위로 읽어옵니다.
    torch_max_allocated gpu mem의 경우 해당 함수 호출 직전까지의 최댓값이며, 해당 함수 호출 시 리셋됩니다.

    Returns:
        Dict[str, int]: _description_
    """
    memory_stats = {
        "cpu_memory(MiB)": get_process_cpu_memory() // 1024**2,
        "gpu_memory(MiB)": get_process_gpu_memory() // 1024**2,
        "gpu_memory(torch_max_allocated, MiB)": torch.cuda.max_memory_allocated() // 1024**2,
    }
    torch.cuda.reset_peak_memory_stats()
    return memory_stats