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 
89
90
91
92
93
94
95
96
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

device instance-attribute 

device = device(device)

enabled instance-attribute 

enabled = all([type == 'cuda', system() == 'Linux', is_available()])

cuda_context_memory instance-attribute 

cuda_context_memory = 0

max_memory_allocated instance-attribute 

max_memory_allocated = 0

max_memory_reserved instance-attribute 

max_memory_reserved = 0

__enter__ 

__enter__()
Source code in SaigeToolkit/util/memory.py 
 98
 99
100
101
102
103
104
105
def __enter__(self):
    if not self.enabled:
        return self
    cuda_empty_cache(self.device)
    torch.cuda.reset_peak_memory_stats(self.device)
    self.memory_already_allocated = torch.cuda.max_memory_allocated(self.device)
    self.memory_already_reserved = torch.cuda.max_memory_reserved(self.device)
    return self

__exit__ 

__exit__(*args, **kwargs)
Source code in SaigeToolkit/util/memory.py 
107
108
109
110
111
112
113
114
115
116
117
118
119
120
def __exit__(self, *args, **kwargs):
    if not self.enabled:
        return
    self.max_memory_allocated = max(
        self.max_memory_allocated,
        torch.cuda.max_memory_allocated(self.device) - self.memory_already_allocated,
    )
    self.max_memory_reserved = max(
        self.max_memory_reserved,
        torch.cuda.max_memory_reserved(self.device) - self.memory_already_reserved,
    )
    self.cuda_context_memory = max(
        0, get_process_gpu_memory() - torch.cuda.memory_reserved(self.device)
    )

start 

start()
Source code in SaigeToolkit/util/memory.py 
122
123
def start(self):
    return self.__enter__()

end 

end(*args, **kwargs)
Source code in SaigeToolkit/util/memory.py 
125
126
def end(self, *args, **kwargs):
    return self.__exit__(*args, **kwargs)

get_results 

get_results() -> dict
Source code in SaigeToolkit/util/memory.py 
128
129
130
131
132
133
134
135
136
def get_results(self) -> dict:
    reserved_and_context = (
        f"{(self.max_memory_reserved + self.cuda_context_memory)/ 1024 ** 2:.0f} MiB"
    )
    return {
        "GPU memory": f"{self.max_memory_allocated / 1024 ** 2:.0f} MiB",
        "GPU memory (+ cache)": f"{self.max_memory_reserved / 1024 ** 2:.0f} MiB",
        "GPU memory (+ cache + CUDA context)": reserved_and_context,
    }

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
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()
    return memory_info.vms if include_virtual else 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 
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
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:
                return 0 if process.usedGpuMemory is None else 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 
54
55
56
57
58
59
60
61
62
63
64
65
66
67
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