| import time |
| from contextlib import contextmanager |
|
|
| import torch |
|
|
|
|
| class TemporalMemoryProfiler: |
| def __init__(self, enabled=False, prefix='temporal_memory'): |
| self.enabled = enabled |
| self.prefix = prefix |
| self.records = {} |
| self.gpu_memory_allocated_mb = 0.0 |
| self.gpu_memory_max_allocated_mb = 0.0 |
|
|
| @contextmanager |
| def measure(self, name): |
| if not self.enabled: |
| yield |
| return |
| self._sync() |
| start = time.perf_counter() |
| try: |
| yield |
| finally: |
| self._sync() |
| elapsed_ms = (time.perf_counter() - start) * 1000.0 |
| self.records[name] = self.records.get(name, 0.0) + elapsed_ms |
| self._capture_memory() |
|
|
| def report(self): |
| if not self.enabled: |
| return {} |
| report = { |
| f'{name}_latency_ms': round(value, 3) |
| for name, value in self.records.items() |
| } |
| report['gpu_memory_allocated_mb'] = round(self.gpu_memory_allocated_mb, 3) |
| report['gpu_memory_max_allocated_mb'] = round(self.gpu_memory_max_allocated_mb, 3) |
| return report |
|
|
| def log(self): |
| report = self.report() |
| if not report: |
| return |
| fields = ' '.join(f'{key}={value:.3f}' for key, value in report.items()) |
| print(f'[{self.prefix}] {fields}') |
|
|
| def reset(self): |
| self.records.clear() |
| self.gpu_memory_allocated_mb = 0.0 |
| self.gpu_memory_max_allocated_mb = 0.0 |
|
|
| def _sync(self): |
| if torch.cuda.is_available(): |
| torch.cuda.synchronize() |
|
|
| def _capture_memory(self): |
| if not torch.cuda.is_available(): |
| return |
| mb = 1024.0 * 1024.0 |
| self.gpu_memory_allocated_mb = torch.cuda.memory_allocated() / mb |
| self.gpu_memory_max_allocated_mb = torch.cuda.max_memory_allocated() / mb |
|
|
