Use binary units for CUDA memory summary (#91854)

milesial · pytorchmergebot · commit 421f40e05143 · 2023-01-14T05:10:51.000Z
To reduce confusion, use for example `KiB` instead of `KB` since we're talking powers of 2 and not 10. https://en.wikipedia.org/wiki/Byte#Multiple-byte_units ``` import torch x = torch.zeros(1024 * 1024, dtype=torch.uint8, device='cuda') print(torch.cuda.memory_summary()) ``` ``` |===========================================================================| | PyTorch CUDA memory summary, device ID 0 | |---------------------------------------------------------------------------| | CUDA OOMs: 0 | cudaMalloc retries: 0 | |===========================================================================| | Metric | Cur Usage | Peak Usage | Tot Alloc | Tot Freed | |---------------------------------------------------------------------------| | Allocated memory | 1024 KiB | 1024 KiB | 1024 KiB | 0 B | | from large pool | 0 KiB | 0 KiB | 0 KiB | 0 B | | from small pool | 1024 KiB | 1024 KiB | 1024 KiB | 0 B | |---------------------------------------------------------------------------| | Active memory | 1024 KiB | 1024 KiB | 1024 KiB | 0 B | | from large pool | 0 KiB | 0 KiB | 0 KiB | 0 B | | from small pool | 1024 KiB | 1024 KiB | 1024 KiB | 0 B | |---------------------------------------------------------------------------| ``` Pull Request resolved: #91854 Approved by: https://github.com/ngimel
diff --git a/torch/cuda/memory.py b/torch/cuda/memory.py
@@ -453,15 +453,15 @@ def memory_summary(device: Union[Device, int] = None, abbreviated: bool = False)
     stats = memory_stats(device=device)
 
     def _format_size(sz, pref_sz):
-        prefixes = ["B ", "KB", "MB", "GB", "TB", "PB"]
+        prefixes = ["B  ", "KiB", "MiB", "GiB", "TiB", "PiB"]
         prefix = prefixes[0]
         for new_prefix in prefixes[1:]:
             if pref_sz < 768 * 1024:
                 break
             prefix = new_prefix
             sz //= 1024
             pref_sz /= 1024
-        return "{:7d} {}".format(sz, prefix)
+        return "{:6d} {}".format(sz, prefix)
 
     def _format_count(cnt, pref_cnt):
         prefixes = [" ", "K", "M"]
