File size: 1,890 Bytes
6a176cb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 | 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
|