| 354 | |
| 355 | |
| 356 | class Timer: |
| 357 | |
| 358 | def __init__( |
| 359 | self, |
| 360 | name, |
| 361 | logger=None, |
| 362 | ): |
| 363 | self.name = name |
| 364 | self.start_time = None |
| 365 | self.end_time = None |
| 366 | self.logger = logger |
| 367 | |
| 368 | @property |
| 369 | def elapsed_time(self): |
| 370 | return self.end_time - self.start_time |
| 371 | |
| 372 | def __enter__(self): |
| 373 | torch.cuda.synchronize() |
| 374 | self.start_time = time() |
| 375 | return self |
| 376 | |
| 377 | def __exit__(self, exc_type, exc_val, exc_tb): |
| 378 | torch.cuda.synchronize() |
| 379 | self.end_time = time() |
| 380 | msg = f'Elapsed time for {self.name}: {self.elapsed_time:.3f} s' |
| 381 | if self.logger is not None: |
| 382 | self.logger.info(msg, main_process_only=True) |
| 383 | else: |
| 384 | print(msg) |
| 385 | |
| 386 | |
| 387 | # https://github.com/hpcaitech/Open-Sora/blob/main/opensora/utils/train_utils.py#L44 |