Spaces:
Sleeping
Sleeping
| import time | |
| import torch | |
| import inspect | |
| def fold_path(fn:str): | |
| ''' Fold a path like `from/to/file.py` to relative `f/t/file.py`. ''' | |
| return '/'.join([p[:1] for p in fn.split('/')[:-1]]) + '/' + fn.split('/')[-1] | |
| def summary_frame_info(frame:inspect.FrameInfo): | |
| ''' Convert a FrameInfo object to a summary string. ''' | |
| return f'{frame.function} @ {fold_path(frame.filename)}:{frame.lineno}' | |
| class GPUMonitor(): | |
| ''' | |
| This monitor is designed for GPU memory analysis. It records the peak memory usage in a period of time. | |
| A snapshot will record the peak memory usage until the snapshot is taken. (After init / reset / previous snapshot.) | |
| ''' | |
| def __init__(self): | |
| self.reset() | |
| self.clear() | |
| self.log_fn = 'gpu_monitor.log' | |
| def snapshot(self, desc:str='snapshot'): | |
| timestamp = time.time() | |
| caller_frame = inspect.stack()[1] | |
| peak_MB = torch.cuda.max_memory_allocated() / 1024 / 1024 | |
| free_mem, total_mem = torch.cuda.mem_get_info(0) | |
| free_mem_MB, total_mem_MB = free_mem / 1024 / 1024, total_mem / 1024 / 1024 | |
| record = { | |
| 'until' : time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(timestamp)), | |
| 'until_raw' : timestamp, | |
| 'position' : summary_frame_info(caller_frame), | |
| 'peak' : peak_MB, | |
| 'peak_msg' : f'{peak_MB:.2f} MB', | |
| 'free' : free_mem_MB, | |
| 'total' : total_mem_MB, | |
| 'free_msg' : f'{free_mem_MB:.2f} MB', | |
| 'total_msg' : f'{total_mem_MB:.2f} MB', | |
| 'desc' : desc, | |
| } | |
| self.max_peak = max(self.max_peak_MB, peak_MB) | |
| self.records.append(record) | |
| self._update_log(record) | |
| self.reset() | |
| return record | |
| def report_latest(self, k:int=1): | |
| import rich | |
| caller_frame = inspect.stack()[1] | |
| caller_info = summary_frame_info(caller_frame) | |
| rich.print(f'{caller_info} -> latest {k} records:') | |
| for rid, record in enumerate(self.records[-k:]): | |
| msg = self._generate_log_msg(record) | |
| rich.print(msg) | |
| def report_all(self): | |
| self.report_latest(len(self.records)) | |
| def reset(self): | |
| torch.cuda.reset_peak_memory_stats() | |
| return | |
| def clear(self): | |
| self.records = [] | |
| self.max_peak_MB = 0 | |
| def _generate_log_msg(self, record): | |
| time = record['until'] | |
| peak = record['peak'] | |
| desc = record['desc'] | |
| position = record['position'] | |
| msg = f'[{time}] ⛰️ {peak:>8.2f} MB 📌 {desc} 🌐 {position}' | |
| return msg | |
| def _update_log(self, record): | |
| msg = self._generate_log_msg(record) | |
| with open(self.log_fn, 'a') as f: | |
| f.write(msg + '\n') | |