| |
| import contextlib |
| import sys |
| import time |
|
|
| import torch |
|
|
| if sys.version_info >= (3, 7): |
|
|
| @contextlib.contextmanager |
| def profile_time(trace_name, |
| name, |
| enabled=True, |
| stream=None, |
| end_stream=None): |
| """Print time spent by CPU and GPU. |
| |
| Useful as a temporary context manager to find sweet spots of code |
| suitable for async implementation. |
| """ |
| if (not enabled) or not torch.cuda.is_available(): |
| yield |
| return |
| stream = stream if stream else torch.cuda.current_stream() |
| end_stream = end_stream if end_stream else stream |
| start = torch.cuda.Event(enable_timing=True) |
| end = torch.cuda.Event(enable_timing=True) |
| stream.record_event(start) |
| try: |
| cpu_start = time.monotonic() |
| yield |
| finally: |
| cpu_end = time.monotonic() |
| end_stream.record_event(end) |
| end.synchronize() |
| cpu_time = (cpu_end - cpu_start) * 1000 |
| gpu_time = start.elapsed_time(end) |
| msg = f'{trace_name} {name} cpu_time {cpu_time:.2f} ms ' |
| msg += f'gpu_time {gpu_time:.2f} ms stream {stream}' |
| print(msg, end_stream) |
|
|
