| 19 | |
| 20 | |
| 21 | class Timer: |
| 22 | timer_map = {} |
| 23 | |
| 24 | def __init__(self, name, enable=False): |
| 25 | if name not in Timer.timer_map: |
| 26 | Timer.timer_map[name] = 0 |
| 27 | self.name = name |
| 28 | self.enable = enable |
| 29 | |
| 30 | def __enter__(self): |
| 31 | if self.enable: |
| 32 | # if torch.cuda.is_available(): |
| 33 | # torch.cuda.synchronize() |
| 34 | self.t = time.time() |
| 35 | |
| 36 | def __exit__(self, exc_type, exc_val, exc_tb): |
| 37 | if self.enable: |
| 38 | # if torch.cuda.is_available(): |
| 39 | # torch.cuda.synchronize() |
| 40 | Timer.timer_map[self.name] += time.time() - self.t |
| 41 | if self.enable: |
| 42 | print(f'[Timer] {self.name}: {Timer.timer_map[self.name]}') |
no outgoing calls
no test coverage detected