| 210 | |
| 211 | |
| 212 | class LoggingCallback(pl.Callback): |
| 213 | @rank_zero_only |
| 214 | def on_validation_end(self, trainer: pl.Trainer, pl_module: pl.LightningModule): |
| 215 | rank_zero_info("***** Validation results *****") |
| 216 | metrics = trainer.callback_metrics |
| 217 | # Log results |
| 218 | for key in sorted(metrics): |
| 219 | if key not in ["log", "progress_bar"]: |
| 220 | rank_zero_info("{} = {}\n".format(key, str(metrics[key]))) |
| 221 | |
| 222 | @rank_zero_only |
| 223 | def on_test_end(self, trainer: pl.Trainer, pl_module: pl.LightningModule): |
| 224 | logger.info("***** Test results *****") |
| 225 | metrics = trainer.callback_metrics |
| 226 | # Log and save results to file |
| 227 | output_test_results_file = os.path.join(pl_module.hparams.output_dir, "test_results.txt") |
| 228 | with open(output_test_results_file, "w") as writer: |
| 229 | for key in sorted(metrics): |
| 230 | if key not in ["log", "progress_bar"]: |
| 231 | logger.info("{} = {}\n".format(key, str(metrics[key]))) |
| 232 | writer.write("{} = {}\n".format(key, str(metrics[key]))) |
| 233 | |
| 234 | |
| 235 | def add_generic_args(parser, root_dir) -> None: |