Simple GPU memory report.
(name)
| 131 | |
| 132 | |
| 133 | def report_memory(name): |
| 134 | """Simple GPU memory report.""" |
| 135 | |
| 136 | mega_bytes = 1024.0 * 1024.0 |
| 137 | string = name + ' memory (MB)' |
| 138 | string += ' | allocated: {}'.format( |
| 139 | torch.cuda.memory_allocated() / mega_bytes) |
| 140 | string += ' | max allocated: {}'.format( |
| 141 | torch.cuda.max_memory_allocated() / mega_bytes) |
| 142 | string += ' | cached: {}'.format(torch.cuda.memory_cached() / mega_bytes) |
| 143 | string += ' | max cached: {}'.format( |
| 144 | torch.cuda.max_memory_cached()/ mega_bytes) |
| 145 | print_rank_0(string) |
| 146 | |
| 147 | |
| 148 | def get_checkpoint_name(checkpoints_path, iteration, release=False, zero=False): |
no test coverage detected