| 86 | |
| 87 | |
| 88 | class SetupCallback(Callback): |
| 89 | def __init__(self, resume, logdir, ckptdir, cfgdir, config): |
| 90 | super().__init__() |
| 91 | self.resume = resume |
| 92 | self.logdir = logdir |
| 93 | self.ckptdir = ckptdir |
| 94 | self.cfgdir = cfgdir |
| 95 | self.config = config |
| 96 | |
| 97 | def on_fit_start(self, trainer, pl_module): |
| 98 | if trainer.global_rank == 0: |
| 99 | # Create logdirs and save configs |
| 100 | os.makedirs(self.logdir, exist_ok=True) |
| 101 | os.makedirs(self.ckptdir, exist_ok=True) |
| 102 | os.makedirs(self.cfgdir, exist_ok=True) |
| 103 | |
| 104 | rank_zero_print("Project config") |
| 105 | rank_zero_print(OmegaConf.to_yaml(self.config)) |
| 106 | OmegaConf.save(self.config, |
| 107 | os.path.join(self.cfgdir, "project.yaml")) |
| 108 | |
| 109 | |
| 110 | class CodeSnapshot(Callback): |
nothing calls this directly
no outgoing calls
no test coverage detected