| 464 | |
| 465 | |
| 466 | class SetupCallback(Callback): |
| 467 | def __init__(self, resume, now, logdir, ckptdir, cfgdir, config, lightning_config): |
| 468 | super().__init__() |
| 469 | self.resume = resume |
| 470 | self.now = now |
| 471 | self.logdir = logdir |
| 472 | self.ckptdir = ckptdir |
| 473 | self.cfgdir = cfgdir |
| 474 | self.config = config |
| 475 | self.lightning_config = lightning_config |
| 476 | |
| 477 | def on_keyboard_interrupt(self, trainer, pl_module): |
| 478 | if trainer.global_rank == 0: |
| 479 | print("Summoning checkpoint.") |
| 480 | ckpt_path = os.path.join(self.ckptdir, "last.ckpt") |
| 481 | trainer.save_checkpoint(ckpt_path) |
| 482 | |
| 483 | def on_pretrain_routine_start(self, trainer, pl_module): |
| 484 | if trainer.global_rank == 0: |
| 485 | # Create logdirs and save configs |
| 486 | os.makedirs(self.logdir, exist_ok=True) |
| 487 | os.makedirs(self.ckptdir, exist_ok=True) |
| 488 | os.makedirs(self.cfgdir, exist_ok=True) |
| 489 | |
| 490 | if "callbacks" in self.lightning_config: |
| 491 | if 'metrics_over_trainsteps_checkpoint' in self.lightning_config['callbacks']: |
| 492 | os.makedirs(os.path.join(self.ckptdir, 'trainstep_checkpoints'), exist_ok=True) |
| 493 | print("Project config") |
| 494 | print(OmegaConf.to_yaml(self.config)) |
| 495 | OmegaConf.save(self.config, |
| 496 | os.path.join(self.cfgdir, "{}-project.yaml".format(self.now))) |
| 497 | |
| 498 | print("Lightning config") |
| 499 | print(OmegaConf.to_yaml(self.lightning_config)) |
| 500 | OmegaConf.save(OmegaConf.create({"lightning": self.lightning_config}), |
| 501 | os.path.join(self.cfgdir, "{}-lightning.yaml".format(self.now))) |
| 502 | |
| 503 | else: |
| 504 | # ModelCheckpoint callback created log directory --- remove it |
| 505 | if not self.resume and os.path.exists(self.logdir): |
| 506 | dst, name = os.path.split(self.logdir) |
| 507 | dst = os.path.join(dst, "child_runs", name) |
| 508 | os.makedirs(os.path.split(dst)[0], exist_ok=True) |
| 509 | try: |
| 510 | os.rename(self.logdir, dst) |
| 511 | except FileNotFoundError: |
| 512 | pass |
| 513 | |
| 514 | |
| 515 | class ImageLogger(Callback): |
nothing calls this directly
no outgoing calls
no test coverage detected