| 60 | |
| 61 | |
| 62 | class gpu_timer: |
| 63 | def __init__(self, beta=0.6) -> None: |
| 64 | self.start_time = None |
| 65 | self.stop_time = None |
| 66 | self.mean_time = None |
| 67 | self.beta = beta |
| 68 | self.first_call = True |
| 69 | |
| 70 | def start(self): |
| 71 | torch.cuda.synchronize() |
| 72 | self.start_time = time.perf_counter() |
| 73 | |
| 74 | def stop(self): |
| 75 | if self.start_time is None: |
| 76 | print("Use start() before stop(). ") |
| 77 | torch.cuda.synchronize() |
| 78 | self.stop_time = time.perf_counter() |
| 79 | elapsed = self.stop_time - self.start_time |
| 80 | self.start_time = None |
| 81 | if self.first_call: |
| 82 | self.mean_time = elapsed |
| 83 | self.first_call = False |
| 84 | else: |
| 85 | self.mean_time = self.beta * self.mean_time + (1 - self.beta) * elapsed |
| 86 | |
| 87 | |
| 88 | def set_seed(seed): |