Simple GPU memory report.
(name)
| 133 | |
| 134 | |
| 135 | def report_memory(name): |
| 136 | """Simple GPU memory report.""" |
| 137 | |
| 138 | mega_bytes = 1024.0 * 1024.0 |
| 139 | string = name + ' memory (MB)' |
| 140 | string += ' | allocated: {}'.format( |
| 141 | torch.cuda.memory_allocated() / mega_bytes) |
| 142 | string += ' | max allocated: {}'.format( |
| 143 | torch.cuda.max_memory_allocated() / mega_bytes) |
| 144 | string += ' | cached: {}'.format(torch.cuda.memory_reserved() / mega_bytes) |
| 145 | string += ' | max cached: {}'.format( |
| 146 | torch.cuda.max_memory_reserved() / mega_bytes) |
| 147 | print_rank0(string) |
| 148 |