| | 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) |
| |
|
