Simple GPU memory report.
(name)
| 177 | |
| 178 | |
| 179 | def report_memory(name): |
| 180 | """Simple GPU memory report.""" |
| 181 | |
| 182 | mega_bytes = 1024.0 * 1024.0 |
| 183 | string = name + ' memory (MB)' |
| 184 | string += ' | allocated: {}'.format( |
| 185 | torch.cuda.memory_allocated() / mega_bytes) |
| 186 | string += ' | max allocated: {}'.format( |
| 187 | torch.cuda.max_memory_allocated() / mega_bytes) |
| 188 | string += ' | cached: {}'.format(torch.cuda.memory_cached() / mega_bytes) |
| 189 | string += ' | max cached: {}'.format( |
| 190 | torch.cuda.memory_reserved() / mega_bytes) |
| 191 | print_rank_0(string) |
| 192 | |
| 193 | |
| 194 | def get_checkpoint_name(checkpoints_path, iteration, release=False, zero=False): |