def _test_loggers_fit_test(tmpdir, logger_class): model = EvalModelTemplate() class StoreHistoryLogger(logger_class): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.history = [] def log_metrics(self, metrics, step): super().log_metrics(metrics, step) self.history.append((step, metrics)) logger_args = _get_logger_args(logger_class, tmpdir) logger = StoreHistoryLogger(**logger_args) if logger_class == WandbLogger: # required mocks for Trainer logger.experiment.id = 'foo' logger.experiment.project_name.return_value = 'bar' if logger_class == CometLogger: logger.experiment.id = 'foo' logger.experiment.project_name = 'bar' if logger_class == TestTubeLogger: logger.experiment.version = 'foo' logger.experiment.name = 'bar' if logger_class == MLFlowLogger: logger = mock_mlflow_run_creation(logger, experiment_id="foo", run_id="bar") trainer = Trainer( max_epochs=1, logger=logger, limit_train_batches=1, limit_val_batches=1, log_every_n_steps=1, default_root_dir=tmpdir, ) trainer.fit(model) trainer.test() log_metric_names = [(s, sorted(m.keys())) for s, m in logger.history] if logger_class == TensorBoardLogger: expected = [ (0, ['hp_metric']), (0, ['epoch', 'train_some_val']), (0, ['early_stop_on', 'epoch', 'val_acc']), (0, ['hp_metric']), (1, ['epoch', 'test_acc', 'test_loss']) ] assert log_metric_names == expected else: expected = [ (0, ['epoch', 'train_some_val']), (0, ['early_stop_on', 'epoch', 'val_acc']), (1, ['epoch', 'test_acc', 'test_loss']) ] assert log_metric_names == expected
def _test_loggers_fit_test(tmpdir, logger_class): class CustomModel(BoringModel): def training_step(self, batch, batch_idx): output = self.layer(batch) loss = self.loss(batch, output) self.log("train_some_val", loss) return {"loss": loss} def validation_epoch_end(self, outputs) -> None: avg_val_loss = torch.stack([x["x"] for x in outputs]).mean() self.log_dict({ "early_stop_on": avg_val_loss, "val_loss": avg_val_loss**0.5 }) def test_epoch_end(self, outputs) -> None: avg_test_loss = torch.stack([x["y"] for x in outputs]).mean() self.log("test_loss", avg_test_loss) class StoreHistoryLogger(logger_class): def __init__(self, *args, **kwargs) -> None: super().__init__(*args, **kwargs) self.history = [] def log_metrics(self, metrics, step): super().log_metrics(metrics, step) self.history.append((step, metrics)) logger_args = _get_logger_args(logger_class, tmpdir) logger = StoreHistoryLogger(**logger_args) if logger_class == WandbLogger: # required mocks for Trainer logger.experiment.id = "foo" logger.experiment.project_name.return_value = "bar" if logger_class == CometLogger: logger.experiment.id = "foo" logger.experiment.project_name = "bar" if logger_class == TestTubeLogger: logger.experiment.version = "foo" logger.experiment.name = "bar" if logger_class == MLFlowLogger: logger = mock_mlflow_run_creation(logger, experiment_id="foo", run_id="bar") model = CustomModel() trainer = Trainer( max_epochs=1, logger=logger, limit_train_batches=1, limit_val_batches=1, log_every_n_steps=1, default_root_dir=tmpdir, ) trainer.fit(model) trainer.test() log_metric_names = [(s, sorted(m.keys())) for s, m in logger.history] if logger_class == TensorBoardLogger: expected = [ (0, ["hp_metric"]), (0, ["epoch", "train_some_val"]), (0, ["early_stop_on", "epoch", "val_loss"]), (0, ["hp_metric"]), (1, ["epoch", "test_loss"]), ] assert log_metric_names == expected else: expected = [ (0, ["epoch", "train_some_val"]), (0, ["early_stop_on", "epoch", "val_loss"]), (1, ["epoch", "test_loss"]), ] assert log_metric_names == expected