| 620 | |
| 621 | |
| 622 | class CUDACallback(Callback): |
| 623 | # see https://github.com/SeanNaren/minGPT/blob/master/mingpt/callback.py |
| 624 | def on_train_epoch_start(self, trainer, pl_module): |
| 625 | # Reset the memory use counter |
| 626 | torch.cuda.reset_peak_memory_stats(trainer.root_gpu) |
| 627 | torch.cuda.synchronize(trainer.root_gpu) |
| 628 | self.start_time = time.time() |
| 629 | |
| 630 | def on_train_epoch_end(self, trainer, pl_module, outputs): |
| 631 | torch.cuda.synchronize(trainer.root_gpu) |
| 632 | max_memory = torch.cuda.max_memory_allocated(trainer.root_gpu) / 2 ** 20 |
| 633 | epoch_time = time.time() - self.start_time |
| 634 | |
| 635 | try: |
| 636 | max_memory = trainer.training_type_plugin.reduce(max_memory) |
| 637 | epoch_time = trainer.training_type_plugin.reduce(epoch_time) |
| 638 | |
| 639 | rank_zero_info(f"Average Epoch time: {epoch_time:.2f} seconds") |
| 640 | rank_zero_info(f"Average Peak memory {max_memory:.2f}MiB") |
| 641 | except AttributeError: |
| 642 | pass |
| 643 | |
| 644 | |
| 645 | if __name__ == "__main__": |
nothing calls this directly
no outgoing calls
no test coverage detected