Simple GPU memory report.
(name)
| 141 | |
| 142 | |
| 143 | def report_memory(name): |
| 144 | """Simple GPU memory report.""" |
| 145 | |
| 146 | mega_bytes = 1024.0 * 1024.0 |
| 147 | string = name + ' memory (MB)' |
| 148 | string += ' | allocated: {}'.format( |
| 149 | torch.cuda.memory_allocated() / mega_bytes) |
| 150 | string += ' | max allocated: {}'.format( |
| 151 | torch.cuda.max_memory_allocated() / mega_bytes) |
| 152 | string += ' | cached: {}'.format(torch.cuda.memory_cached() / mega_bytes) |
| 153 | string += ' | max cached: {}'.format( |
| 154 | torch.cuda.memory_reserved() / mega_bytes) |
| 155 | print_rank_0(string) |
| 156 | |
| 157 | |
| 158 | def get_checkpoint_name(checkpoints_path, iteration, release=False, zero=False): |