diff --git a/pytorch_lightning/trainer/callback_config.py b/pytorch_lightning/trainer/callback_config.py index c3334ab5..7d0496c8 100644 --- a/pytorch_lightning/trainer/callback_config.py +++ b/pytorch_lightning/trainer/callback_config.py @@ -2,7 +2,6 @@ import os from abc import ABC from pytorch_lightning.callbacks import ModelCheckpoint, EarlyStopping -from pytorch_lightning.logging import TensorBoardLogger class TrainerCallbackConfigMixin(ABC): @@ -50,7 +49,7 @@ class TrainerCallbackConfigMixin(ABC): if self.weights_save_path is None: self.weights_save_path = self.default_save_path - def configure_early_stopping(self, early_stop_callback, logger): + def configure_early_stopping(self, early_stop_callback): if early_stop_callback is True: self.early_stop_callback = EarlyStopping( monitor='val_loss', @@ -75,18 +74,3 @@ class TrainerCallbackConfigMixin(ABC): else: self.early_stop_callback = early_stop_callback self.enable_early_stop = True - - # configure logger - if logger is True: - # default logger - self.logger = TensorBoardLogger( - save_dir=self.default_save_path, - version=self.slurm_job_id, - name='lightning_logs' - ) - self.logger.rank = 0 - elif logger is False: - self.logger = None - else: - self.logger = logger - self.logger.rank = 0 diff --git a/pytorch_lightning/trainer/logging.py b/pytorch_lightning/trainer/logging.py index 7c445c07..adc6ec5d 100644 --- a/pytorch_lightning/trainer/logging.py +++ b/pytorch_lightning/trainer/logging.py @@ -3,6 +3,7 @@ from abc import ABC import torch from pytorch_lightning.core import memory +from pytorch_lightning.logging import TensorBoardLogger class TrainerLoggingMixin(ABC): @@ -21,6 +22,21 @@ class TrainerLoggingMixin(ABC): self.use_ddp2 = None self.num_gpus = None + def configure_logger(self, logger): + if logger is True: + # default logger + self.logger = TensorBoardLogger( + save_dir=self.default_save_path, + version=self.slurm_job_id, + name='lightning_logs' + ) + self.logger.rank = 0 + elif logger is False: + self.logger = None + else: + self.logger = logger + self.logger.rank = 0 + def log_metrics(self, metrics, grad_norm_dic, step=None): """Logs the metric dict passed in. diff --git a/pytorch_lightning/trainer/trainer.py b/pytorch_lightning/trainer/trainer.py index 0e31d5b2..9cc9b559 100644 --- a/pytorch_lightning/trainer/trainer.py +++ b/pytorch_lightning/trainer/trainer.py @@ -558,10 +558,12 @@ class Trainer(TrainerIOMixin, self.current_epoch = 0 self.total_batches = 0 + # configure logger + self.configure_logger(logger) + # configure early stop callback # creates a default one if none passed in - self.early_stop_callback = None - self.configure_early_stopping(early_stop_callback, logger) + self.configure_early_stopping(early_stop_callback) self.reduce_lr_on_plateau_scheduler = None