| 597 | |
| 598 | |
| 599 | class CUDACallback(Callback): |
| 600 | # see https://github.com/SeanNaren/minGPT/blob/master/mingpt/callback.py |
| 601 | def on_train_epoch_start(self, trainer, pl_module): |
| 602 | # Reset the memory use counter |
| 603 | torch.cuda.reset_peak_memory_stats(trainer.root_gpu) |
| 604 | torch.cuda.synchronize(trainer.root_gpu) |
| 605 | self.start_time = time.time() |
| 606 | |
| 607 | def on_train_epoch_end(self, trainer, pl_module, outputs): |
| 608 | torch.cuda.synchronize(trainer.root_gpu) |
| 609 | max_memory = torch.cuda.max_memory_allocated(trainer.root_gpu) / 2 ** 20 |
| 610 | epoch_time = time.time() - self.start_time |
| 611 | |
| 612 | try: |
| 613 | max_memory = trainer.training_type_plugin.reduce(max_memory) |
| 614 | epoch_time = trainer.training_type_plugin.reduce(epoch_time) |
| 615 | |
| 616 | rank_zero_info(f"Average Epoch time: {epoch_time:.2f} seconds") |
| 617 | rank_zero_info(f"Average Peak memory {max_memory:.2f}MiB") |
| 618 | except AttributeError: |
| 619 | pass |
| 620 | |
| 621 | |
| 622 | if __name__ == "__main__": |
nothing calls this directly
no outgoing calls
no test coverage detected