Print memory usage
(title: str, gpu_id: int, device_id: int | None = None)
| 703 | |
| 704 | |
| 705 | def print_gpu_memory_use(title: str, gpu_id: int, device_id: int | None = None) -> None: |
| 706 | """Print memory usage""" |
| 707 | import pynvml |
| 708 | |
| 709 | if device_id is None: |
| 710 | device_id = gpu_id |
| 711 | |
| 712 | pynvml.nvmlInit() |
| 713 | handle = pynvml.nvmlDeviceGetHandleByIndex(device_id) |
| 714 | meminfo = pynvml.nvmlDeviceGetMemoryInfo(handle) |
| 715 | pynvml.nvmlShutdown() |
| 716 | |
| 717 | paddle_max_reserved = paddle.device.cuda.max_memory_reserved(gpu_id) |
| 718 | paddle_max_allocated = paddle.device.cuda.max_memory_allocated(gpu_id) |
| 719 | paddle_reserved = paddle.device.cuda.memory_reserved(gpu_id) |
| 720 | paddle_allocated = paddle.device.cuda.memory_allocated(gpu_id) |
| 721 | |
| 722 | print( |
| 723 | f"\n{title}:", |
| 724 | f"\n\tDevice Total memory(GiB): {meminfo.total / 1024.0 / 1024.0 / 1024.0}", |
| 725 | f"\n\tDevice Used memory(GiB): {meminfo.used / 1024.0 / 1024.0 / 1024.0}", |
| 726 | f"\n\tDevice Free memory(GiB): {meminfo.free / 1024.0 / 1024.0 / 1024.0}", |
| 727 | f"\n\tPaddle max memory Reserved(GiB): {paddle_max_reserved / 1024.0 / 1024.0 / 1024.0}", |
| 728 | f"\n\tPaddle max memory Allocated(GiB): {paddle_max_allocated / 1024.0 / 1024.0 / 1024.0}", |
| 729 | f"\n\tPaddle memory Reserved(GiB): {paddle_reserved / 1024.0 / 1024.0 / 1024.0}", |
| 730 | f"\n\tPaddle memory Allocated(GiB): {paddle_allocated / 1024.0 / 1024.0 / 1024.0}\n", |
| 731 | ) |
| 732 | |
| 733 | |
| 734 | def ceil_div(x: int, y: int) -> int: |